Fix NaNs in TensorBoard

NVlabs · Oct 21, 2020 · 6f3942c · 6f3942c
1 parent a831de2
commit 6f3942c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dnnlib/tflib/optimizer.py b/dnnlib/tflib/optimizer.py
@@ -132,7 +132,7 @@ def register_gradients(self, loss: TfExpression, trainable_vars: Union[List, dic
         assert all(var.shape.as_list() == var_shape for var, var_shape in zip(trainable_vars, self._gradient_shapes))
 
         # Report memory usage if requested.
-        deps = []
+        deps = [loss]
         if self._report_mem_usage:
             self._report_mem_usage = False
             try: