To be technically correct, need to check the in-place _ ver of op

huggingface · Nov 28, 2024 · e3e434b · e3e434b
1 parent 7c32d3b
commit e3e434b
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/timm/optim/lion.py b/timm/optim/lion.py
@@ -143,7 +143,7 @@ def lion(
     if foreach is None:
         try:
             # cannot do foreach if this overload doesn't exist when caution enabled
-            foreach = not caution or 'Scalar' in torch.ops.aten._foreach_maximum.overloads()
+            foreach = not caution or 'Scalar' in torch.ops.aten._foreach_maximum_.overloads()
         except:
             foreach = False
 

diff --git a/timm/optim/nadamw.py b/timm/optim/nadamw.py
@@ -171,7 +171,7 @@ def nadamw(
     if foreach is None:
         try:
             # cannot do foreach if this overload doesn't exist when caution enabled
-            foreach = not caution or 'Scalar' in torch.ops.aten._foreach_maximum.overloads()
+            foreach = not caution or 'Scalar' in torch.ops.aten._foreach_maximum_.overloads()
         except:
             foreach = False