Restored functional trainer

1 files changed, 5 insertions, 3 deletions

diff --git a/training/util.py b/training/util.py
index f46cc61..557b196 100644
--- a/training/util.py
+++ b/training/util.py
@@ -180,11 +180,13 @@ class EMAModel:
 
     @contextmanager
     def apply_temporary(self, parameters: Iterable[torch.nn.Parameter]):
+        parameters = list(parameters)
+        original_params = [p.clone() for p in parameters]
+        self.copy_to(parameters)
+
         try:
-            parameters = list(parameters)
-            original_params = [p.clone() for p in parameters]
-            self.copy_to(parameters)
             yield
         finally:
             for o_param, param in zip(original_params, parameters):
                 param.data.copy_(o_param.data)
+            del original_params