Version bump and only store params requiring grad.v0.2

1 files changed, 3 insertions, 1 deletions

diff --git a/torch_ema/ema.py b/torch_ema/ema.py
index 447fd1e..7771ef7 100644
--- a/torch_ema/ema.py
+++ b/torch_ema/ema.py
@@ -67,7 +67,9 @@ class ExponentialMovingAverage:
           parameters: Iterable of `torch.nn.Parameter`; the parameters to be
             temporarily stored.
         """
-        self.collected_params = [param.clone() for param in parameters]
+        self.collected_params = [param.clone()
+                                 for param in parameters
+                                 if param.requires_grad]
 
     def restore(self, parameters):
         """