make no-grad-sync yield None

axonn-ai · Jul 1, 2024 · de88509 · de88509
1 parent 57a895d
commit de88509
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/axonn/intra_layer/__init__.py b/axonn/intra_layer/__init__.py
@@ -196,6 +196,7 @@ def no_grad_sync():
     old_val = NO_GRADIENT_SYNC
     try:
         NO_GRADIENT_SYNC = True
+        yield None
     finally:
         NO_GRADIENT_SYNC = old_val