def call(self, inputs, training=None, **kwargs):
if training and self.p > 0:
masks = [tf_bernoulli(tf.shape(x)[:2], 1 - self.p)
for x in inputs]
total = sum(masks)
scale = len(inputs) / tf.reduce_max(tf.ones_like(total))
masks = [mask * scale for mask in masks]
inputs = [item * tf.expand_dims(mask, axis=-1)
for item, mask in zip(inputs, masks)]
return inputs
def get_mask(x, p):
mask = tf_bernoulli(tf.shape(x), 1 - p)
mask = mask / (1 - p)
return mask
