title: How do I make a multi-output Tensorflow 2 custom training loop for both regression and classification?
body: <p>I made a minimally reproducible example with the Iris dataset. I made an entire neural network that predicts the last column of the Iris features. I also want to output the target (category). So, the network must minimize two different loss functions (continuous, and categorical). All is set for the continuous target in the next example. But, how do I turn it into a multi-output problem?</p>
  
  <pre><code>import tensorflow as tf
  from tensorflow.keras.layers import Dense
  from tensorflow.keras import Model
  from sklearn.datasets import load_iris
  tf.keras.backend.set_floatx('float64')
  iris, target = load_iris(return_X_y=True)
  
  X = iris[:, :3]
  y = iris[:, 3]
  z = target
  
  ds = tf.data.Dataset.from_tensor_slices((X, y, z)).batch(8)
  
  class MyModel(Model):
      def __init__(self):
          super(MyModel, self).__init__()
          self.d0 = Dense(16, activation='relu')
          self.d1 = Dense(32, activation='relu')
          self.d2 = Dense(1)
  
      def call(self, x):
          x = self.d0(x)
          x = self.d1(x)
          x = self.d2(x)
          return x
  
  model = MyModel()
  
  loss_object = tf.keras.losses.MeanAbsoluteError()
  optimizer = tf.keras.optimizers.Adam(learning_rate=1e-4)
  
  loss = tf.keras.metrics.Mean(name='categorical loss')
  error = tf.keras.metrics.MeanAbsoluteError()
  
  @tf.function
  def train_step(inputs, target):
      with tf.GradientTape() as tape:
          output = model(inputs)
          run_loss = loss_object(target, output)
  
      gradients = tape.gradient(run_loss, model.trainable_variables)
      optimizer.apply_gradients(zip(gradients, model.trainable_variables))
  
      loss(run_loss)
      error(target, output)
  
  
  for epoch in range(50):
      for xx, yy, zz in ds: # what to do with zz, the categorical target?
          train_step(xx, yy)
  
      template = 'Epoch {:&gt;2}, MAE: {:&gt;5.2f}'
      print(template.format(epoch+1,
                          loss.result()))
  
      loss.reset_states()
      error.reset_states()
  </code></pre>