From de88509a4eaf10ce9f9771cde725d0ccf11a9e68 Mon Sep 17 00:00:00 2001 From: Siddharth Singh Date: Mon, 1 Jul 2024 04:13:12 -0700 Subject: [PATCH] make no-grad-sync yield None --- axonn/intra_layer/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/axonn/intra_layer/__init__.py b/axonn/intra_layer/__init__.py index 98a4546..4700645 100644 --- 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