Use trick shown here https://pytorch.org/tutorials/recipes/recipes/tu…

…ning_guide.html#use-parameter-grad-none-instead-of-model-zero-grad-or-optimizer-zero-grad

ptgnn/baseneuralmodel/trainer.py

@@ -214,7 +214,9 @@ def _run_training(
                     parallelize=parallelize,
                 )
             ):
-                optimizer.zero_grad()
+                optimizer.zero_grad(
+                    set_to_none=True
+                )  # https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html#torch.optim.Optimizer.zero_grad)
                 with catch_cuda_oom(self._catch_cuda_ooms):
                     with torch.cuda.amp.autocast(enabled=self._enable_amp):
                         mb_loss = self.neural_module(**mb_data)