From 632ce00b54ffeacfc18f44f10827f167ab3ac37c Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Sun, 15 Jan 2023 21:06:16 +0100
Subject: Restored functional trainer

---
 training/util.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

(limited to 'training/util.py')

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
-- 
cgit v1.2.3-54-g00ecf