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_llama_pretrain.py
+++ b/tests/e2e/test_llama_pretrain.py
@@ -62,7 +62,7 @@ class TestPretrainLlama:
                "val_set_size": 0.0,
                "output_dir": temp_dir,
                "learning_rate": 0.00001,
-                "optimizer": "adamw_torch",
+                "optimizer": "adamw_torch_fused",
                "lr_scheduler": "cosine",
                "save_safetensors": True,
                "bf16": "auto",