fix resnet examples (#7899)

pytorch · Aug 22, 2024 · df1f994 · df1f994
1 parent ae28308
commit df1f994
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/examples/train_resnet_base.py b/examples/train_resnet_base.py
@@ -48,6 +48,7 @@ def step_fn(self, data, target):
     loss = self.loss_fn(output, target)
     loss.backward()
     self.run_optimizer()
+    return loss
 
   def train_loop_fn(self, loader, epoch):
     tracker = xm.RateTracker()