Update test_thunder_ddp.py

Lightning-AI · May 23, 2024 · 5a09b5a · 5a09b5a
1 parent a56652d
commit 5a09b5a
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/test_thunder_ddp.py b/tests/test_thunder_ddp.py
@@ -44,6 +44,8 @@ def test_no_backward_sync(choice):
 
         cmodel = model = thunder.jit(model)
     model = fabric.setup(model)
+    if "thunder" not in choice:
+        cmodel = model
 
     # 6 iters, 3 grad accumulation iters
     for i, enabled in enumerate((True, True, False, True, True, False), 1):