mark step before

pytorch · May 17, 2024 · 7a18a76 · 7a18a76
1 parent ae0f791
commit 7a18a76
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/torch_xla/torch_xla.py b/torch_xla/torch_xla.py
@@ -62,6 +62,9 @@ def step():
   works with `xla.step` but does not follow best practices will become errors in
   future releases. See https://github.com/pytorch/xla/issues/6751 for context.
   """
+  # Clear pending operations
+  xm.mark_step()
+
   try:
     yield
   finally: