Fix Pytorch 2.0 breakage for Lookahead optimizer adapter

reinhud · Jun 2, 2023 · 700aebc · 700aebc
1 parent cd950e6
commit 700aebc
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/timm/optim/lookahead.py b/timm/optim/lookahead.py
@@ -4,6 +4,9 @@
 
 Hacked together by / Copyright 2020 Ross Wightman
 """
+from collections import OrderedDict
+from typing import Callable, Dict
+
 import torch
 from torch.optim.optimizer import Optimizer
 from collections import defaultdict
@@ -12,6 +15,8 @@
 class Lookahead(Optimizer):
     def __init__(self, base_optimizer, alpha=0.5, k=6):
         # NOTE super().__init__() not called on purpose
+        self._optimizer_step_pre_hooks: Dict[int, Callable] = OrderedDict()
+        self._optimizer_step_post_hooks: Dict[int, Callable] = OrderedDict()
         if not 0.0 <= alpha <= 1.0:
             raise ValueError(f'Invalid slow update rate: {alpha}')
         if not 1 <= k: