code: def mask_loss(input_tensor, binary_tensor):
      
      with tf.variable_scope("mask_loss"):
          mask = tf.cast(tf.cast(binary_tensor, tf.bool), tf.float32)
  
          return input_tensor * mask
docs: Mask a loss by using a tensor filled with 0 or 1.
  
      :param input_tensor: A float tensor of shape [batch_size, ...] representing the loss/cross_entropy
      :param binary_tensor: A float tensor of shape [batch_size, ...] representing the mask.
      :return: A float tensor of shape [batch_size, ...] representing the masked loss.