Internal Code Change

PiperOrigin-RevId: 590739394
tensorflow · Dec 13, 2023 · 90fdc49 · 90fdc49
1 parent 1b0203e
commit 90fdc49
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tensorflow_graphics/projects/points_to_3Dobjects/models/singleobjectmodel.py b/tensorflow_graphics/projects/points_to_3Dobjects/models/singleobjectmodel.py
@@ -106,10 +106,10 @@ def train_sample(self, sample, optimizer, global_batch_size):
     self.decode_batch(sample, inputs)
 
     with tf.GradientTape() as t:
-      predictions = self.network(inputs['image'])  # pytype: disable=key-error
+      predictions = self.network(inputs['image'])
       outputs = self.decode_predictions(predictions, inputs, {})
       outputs = self.compute_loss(inputs, outputs, global_batch_size)
-      network_gradients = t.gradient(outputs['loss/total'],  # pytype: disable=key-error
+      network_gradients = t.gradient(outputs['loss/total'],
                                      self.network.trainable_weights)
     optimizer.apply_gradients(zip(network_gradients,
                                   self.network.trainable_weights))