From 9b27f848760ceff70ee6a70771be6a137e11790b Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Thu, 28 Nov 2024 13:46:17 -0800 Subject: [PATCH] To be technically correct, need to check the in-place _ ver of op --- timm/optim/lion.py | 2 +- timm/optim/nadamw.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/timm/optim/lion.py b/timm/optim/lion.py index 980a0713..18607232 100644 --- 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 index 17eb6fd0..d9933026 100644 --- 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