Fix bug in test to enable kernel

e9898a60 · Christina Floristean · a4fbd978 · e9898a60
Commit e9898a60 authored Sep 14, 2023 by Christina Floristean
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

tests/test_deepspeed_evo_attention.py tests/test_deepspeed_evo_attention.py +1 -1

No files found.
--- a/tests/test_deepspeed_evo_attention.py
+++ b/tests/test_deepspeed_evo_attention.py
@@ -160,7 +160,7 @@ class TestDeepSpeedKernel(unittest.TestCase):
            out_repro = model(batch)

            # Enable kernel
-            model.globals.use_deepspeed_evo_attention = False
+            model.globals.use_deepspeed_evo_attention = True
            out_repro_ds = model(batch)

            out_repro = tensor_tree_map(lambda t: t.cpu(), out_repro)