title: Model with multiple outputs and custom loss function
body: <p>I'm trying to train a model that has multiple outputs and a custom loss function using keras, but I'm getting some error <code>tensorflow.python.framework.errors_impl.OperatorNotAllowedInGraphError: iterating over ``tf.Tensor`` is not allowed in Graph execution. Use Eager execution or decorate this function with @tf.function.</code></p>
  
  <p>It's hard to debug it because I'm doing <code>model.compile</code> and <code>model.fit</code>. I think it has something to do with how models are supposed to be defined when having multiple outputs, but I can't find good documentation on this. The guide specifies how to have models with multiple outputs suing the functional API, and has an example for this, but it doesn't clarify how custom loss functions should work when subclassing the <code>Model</code> API. My code is as follows:</p>
  
  <pre><code>class DeepEnsembles(Model):
  
      def __init__(self, **kwargs):
          super(DeepEnsembles, self).__init__()
  
          self.num_models = kwargs.get('num_models')
          model = kwargs.get('model')
  
          self.mean = [model(**dict(**kwargs)) for _ in range(self.num_models)]
  
          self.variance = [model(**dict(**kwargs)) for _ in range(self.num_models)]
  
      def call(self, inputs, training=None, mask=None):
          mean_predictions = []
          variance_predictions = []
          for idx in range(self.num_models):
              mean_predictions.append(self.mean[idx](inputs, training=training))
              variance_predictions.append(self.variance[idx](inputs, training=training))
          mean_stack = tf.stack(mean_predictions)
          variance_stack = tf.stack(variance_predictions)
  
          return mean_stack, variance_stack
  </code></pre>
  
  <p>And where MLP is the following:</p>
  
  <pre><code>class MLP(Model):
      def __init__(self, **kwargs):
          super(MLP, self).__init__()
  
          # Initialization parameters
          self.num_inputs = kwargs.get('num_inputs', 779)
          self.num_outputs = kwargs.get('num_outputs', 1)
          self.hidden_size = kwargs.get('hidden_size', 256)
          self.activation = kwargs.get('activation', 'relu')
  
          # Optional parameters
          self.p = kwargs.get('p', 0.05)
  
          self.model = tf.keras.Sequential([
              layers.Dense(self.hidden_size, activation=self.activation, input_shape=(self.num_inputs,)),
              layers.Dropout(self.p),
              layers.Dense(self.hidden_size, activation=self.activation),
              layers.Dropout(self.p),
              layers.Dense(self.num_outputs)
           ])
  
      def call(self, inputs, training=None, mask=None):
          output = self.model(inputs, training=training)
          return output
  
  </code></pre>
  
  <p>I'm trying to minimize a custom loss function </p>
  
  <pre><code>class GaussianNLL(Loss):
  
      def __init__(self):
          super(GaussianNLL, self).__init__()
  
      def call(self, y_true, y_pred):
  
          mean, variance = y_pred
          variance = variance + 0.0001
          nll = (tf.math.log(variance) / 2 + ((y_true - mean) ** 2) / (2 * variance))
          nll = tf.math.reduce_mean(nll)
          return nll
  
  </code></pre>
  
  <p>Finally, this is how I try to train it:</p>
  
  <pre><code>    ensembles_params = {'num_models': 5, 'model': MLP, 'p': 0}
      model = DeepEnsembles(**ensembles_params)
      loss_fn = GaussianNLL()
      optimizer = tf.keras.optimizers.Adam(learning_rate=1e-4)
      epochs = 10000
  
      model.compile(optimizer='adam',
                    loss=loss_fn,
                    metrics=['mse', 'mae'])
      history = model.fit(x_train, y_train,
                          batch_size=2048,
                          epochs=10000,
                          verbose=0,
                          validation_data=(x_val, y_val))
  </code></pre>
  
  <p>Which results in the above error. Any pointers? In particular, the whole stack trace is </p>
  
  <pre><code>Traceback (most recent call last):
    File "/home/emilio/anaconda3/lib/python3.7/contextlib.py", line 130, in __exit__
      self.gen.throw(type, value, traceback)
    File "/home/emilio/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/ops/variable_scope.py", line 2803, in variable_creator_scope
      yield
    File "/home/emilio/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training_v2.py", line 235, in fit
      use_multiprocessing=use_multiprocessing)
    File "/home/emilio/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training_v2.py", line 593, in _process_training_inputs
      use_multiprocessing=use_multiprocessing)
    File "/home/emilio/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training_v2.py", line 646, in _process_inputs
      x, y, sample_weight=sample_weights)
    File "/home/emilio/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training.py", line 2360, in _standardize_user_data
      self._compile_from_inputs(all_inputs, y_input, x, y)
    File "/home/emilio/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training.py", line 2618, in _compile_from_inputs
      experimental_run_tf_function=self._experimental_run_tf_function)
    File "/home/emilio/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/training/tracking/base.py", line 457, in _method_wrapper
      result = method(self, *args, **kwargs)
    File "/home/emilio/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training.py", line 446, in compile
      self._compile_weights_loss_and_weighted_metrics()
    File "/home/emilio/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/training/tracking/base.py", line 457, in _method_wrapper
      result = method(self, *args, **kwargs)
    File "/home/emilio/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training.py", line 1592, in _compile_weights_loss_and_weighted_metrics
      self.total_loss = self._prepare_total_loss(masks)
    File "/home/emilio/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training.py", line 1652, in _prepare_total_loss
      per_sample_losses = loss_fn.call(y_true, y_pred)
    File "/home/emilio/fault_detection/tensorflow_code/tf_utils/loss.py", line 13, in call
      mean, variance = y_pred
    File "/home/emilio/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py", line 539, in __iter__
      self._disallow_iteration()
    File "/home/emilio/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py", line 535, in _disallow_iteration
      self._disallow_in_graph_mode("iterating over `tf.Tensor`")
    File "/home/emilio/anaconda3/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py", line 515, in _disallow_in_graph_mode
      " this function with @tf.function.".format(task))
  tensorflow.python.framework.errors_impl.OperatorNotAllowedInGraphError: iterating over `tf.Tensor` is not allowed in Graph execution. Use Eager execution or decorate this function with @tf.function.
  
  
  </code></pre>
  
  <p>So it's clearly related to the loss function. But the model's forward pass outputs a tuple, which I unpack in the loss function, so I don't know why is this an issue.</p>