Skip to content

Commit 835a1a6

Browse files
committed
Add guard around 'somewhat' newer torch RAdam / NAdam imports
1 parent 191755f commit 835a1a6

File tree

1 file changed

+7
-2
lines changed

1 file changed

+7
-2
lines changed

timm/optim/__init__.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,13 @@
2121
from .sgdp import SGDP
2222
from .sgdw import SGDW
2323

24-
# bring torch optim into timm.optim namespace for consistency
25-
from torch.optim import Adadelta, Adagrad, Adamax, Adam, NAdam, RAdam, RMSprop, SGD
24+
# bring common torch.optim Optimizers into timm.optim namespace for consistency
25+
from torch.optim import Adadelta, Adagrad, Adamax, Adam, AdamW, RMSprop, SGD
26+
try:
27+
# in case any very old torch versions being used
28+
from torch.optim import NAdam, RAdam
29+
except ImportError:
30+
pass
2631

2732
from ._optim_factory import list_optimizers, get_optimizer_class, get_optimizer_info, OptimInfo, OptimizerRegistry, \
2833
create_optimizer_v2, create_optimizer, optimizer_kwargs

0 commit comments

Comments
 (0)