fix(solver): fix gpu fused adamw condition (#196)

InternLM · Apr 11, 2024 · 616f4db · 616f4db
1 parent 7a1dcfd
commit 616f4db
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/internlm/model/ops/fusion_ops_import_helper.py b/internlm/model/ops/fusion_ops_import_helper.py
@@ -138,7 +138,8 @@ def try_import_FusedAdamW():
     backend = internlm_accelerator.get_accelerator_backend()
     try:
         if backend is AcceleratorType.GPU:
-            adam_extra_kwargs["fused"] = True
+            if torch.__version__ >= "2.1.0":
+                adam_extra_kwargs["fused"] = True
 
             if gpc.is_rank_for_log():
                 logger.warning(