model must be deepcopyable to use the feature

lucidrains · Jun 20, 2022 · 2c47f5d · 2c47f5d
1 parent b31f3b4
commit 2c47f5d
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/ema_pytorch/ema_pytorch.py b/ema_pytorch/ema_pytorch.py
@@ -35,7 +35,12 @@ def __init__(
         super().__init__()
         self.beta = beta
         self.online_model = model
-        self.ema_model = copy.deepcopy(model)
+
+        try:
+            self.ema_model = copy.deepcopy(model)
+        except:
+            print('Your model was not copyable. Please make sure you are not using any LazyLinear')
+            exit()
 
         self.update_every = update_every
         self.update_after_step = update_after_step