fix optimizer zero_grad

apache · Jul 18, 2020 · ce535fc · ce535fc
1 parent c9b9b32
commit ce535fc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/submarine-sdk/pysubmarine/submarine/ml/pytorch/model/base_pytorch_model.py b/submarine-sdk/pysubmarine/submarine/ml/pytorch/model/base_pytorch_model.py
@@ -72,8 +72,8 @@ def train(self, train_loader):
             sample, target = batch
             output = self.model(sample)
             loss = self.loss(output, target)
-            loss.backward()
             self.optimizer.zero_grad()
+            loss.backward()
             self.optimizer.step()
 
     def evaluate(self):