1 files changed, 1 insertions, 1 deletions

diff --git a/training/util.py b/training/util.py
index 6f1e85a..bed7111 100644
--- a/training/util.py
+++ b/training/util.py
@@ -262,7 +262,7 @@ class EMAModel:
                 raise ValueError("collected_params and shadow_params must have the same length")
 
     @contextmanager
-    def apply_temporary(self, parameters):
+    def apply_temporary(self, parameters: Iterable[torch.nn.Parameter]):
         try:
             parameters = list(parameters)
             original_params = [p.clone() for p in parameters]