title: Combining gradients from different "networks" in TensorFlow2
body: <p>I'm trying to combine a few &quot;networks&quot; into one final loss function. I'm wondering if what I'm doing is &quot;legal&quot;, as of now I can't seem to make this work. I'm using tensorflow probability :</p>
  <p>The main problem is here:</p>
  <pre><code># Get gradients of the loss wrt the weights.
  gradients = tape.gradient(loss, [m_phis.trainable_weights, m_mus.trainable_weights, m_sigmas.trainable_weights])
  
  # Update the weights of our linear layer.
  optimizer.apply_gradients(zip(gradients, [m_phis.trainable_weights, m_mus.trainable_weights, m_sigmas.trainable_weights])
  </code></pre>
  <p>Which gives me None gradients and throws on apply gradients:</p>
  <blockquote>
  <p>AttributeError: 'list' object has no attribute 'device'</p>
  </blockquote>
  <p>Full code:</p>
  <pre><code>univariate_gmm = tfp.distributions.MixtureSameFamily(
          mixture_distribution=tfp.distributions.Categorical(probs=phis_true),
          components_distribution=tfp.distributions.Normal(loc=mus_true,scale=sigmas_true)
      )
  x = univariate_gmm.sample(n_samples, seed=random_seed).numpy()
  dataset = tf.data.Dataset.from_tensor_slices(x) 
  dataset = dataset.shuffle(buffer_size=1024).batch(64)  
  
  m_phis = keras.layers.Dense(2, activation=tf.nn.softmax)
  m_mus = keras.layers.Dense(2)
  m_sigmas = keras.layers.Dense(2, activation=tf.nn.softplus)
  
  def neg_log_likelihood(y, phis, mus, sigmas):
      a = tfp.distributions.Normal(loc=mus[0],scale=sigmas[0]).prob(y)
      b = tfp.distributions.Normal(loc=mus[1],scale=sigmas[1]).prob(y)
      c = np.log(phis[0]*a + phis[1]*b)
      return tf.reduce_sum(-c, axis=-1)
  
  # Instantiate a logistic loss function that expects integer targets.
  loss_fn = neg_log_likelihood
  
  # Instantiate an optimizer.
  optimizer = tf.keras.optimizers.SGD(learning_rate=1e-3)
  
  # Iterate over the batches of the dataset.
  for step, y in enumerate(dataset):
      
      yy = np.expand_dims(y, axis=1)
  
      # Open a GradientTape.
      with tf.GradientTape() as tape:
          
          # Forward pass.
          phis = m_phis(yy)
          mus = m_mus(yy)
          sigmas = m_sigmas(yy)
  
          # Loss value for this batch.
          loss = loss_fn(yy, phis, mus, sigmas)
  
      # Get gradients of the loss wrt the weights.
      gradients = tape.gradient(loss, [m_phis.trainable_weights, m_mus.trainable_weights, m_sigmas.trainable_weights])
  
      # Update the weights of our linear layer.
      optimizer.apply_gradients(zip(gradients, [m_phis.trainable_weights, m_mus.trainable_weights, m_sigmas.trainable_weights]))
  
      # Logging.
      if step % 100 == 0:
          print(&quot;Step:&quot;, step, &quot;Loss:&quot;, float(loss))
  </code></pre>