updating to fused (#2293)

2025-01-30 16:45:56 +00:00
parent 8779997ba5
commit ac471a697a
37 changed files with 51 additions and 51 deletions
--- a/tests/e2e/test_packing_loss.py
+++ b/tests/e2e/test_packing_loss.py
@@ -49,7 +49,7 @@ class TestPackedLlama(unittest.TestCase):
                "gradient_accumulation_steps": 4,
                "output_dir": temp_dir,
                "learning_rate": 0.00001,
-                "optimizer": "adamw_torch",
+                "optimizer": "adamw_torch_fused",
                "lr_scheduler": "cosine",
                "max_steps": 5,
                "use_tensorboard": True,