From df1f994e830bbf40e473e3509dd2c27e819c93f1 Mon Sep 17 00:00:00 2001
From: JackCaoG <59073027+JackCaoG@users.noreply.github.com>
Date: Thu, 22 Aug 2024 10:14:08 -0700
Subject: [PATCH] fix resnet examples (#7899)

---
 examples/train_resnet_base.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/examples/train_resnet_base.py b/examples/train_resnet_base.py
index daab0b39ae5..c4a8890e9be 100644
--- 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()