Filter by requires_grad before zip() in update().

1 files changed, 2 insertions, 2 deletions

diff --git a/torch_ema/ema.py b/torch_ema/ema.py
index 32ed7ca..4e9dd99 100644
--- a/torch_ema/ema.py
+++ b/torch_ema/ema.py
@@ -42,9 +42,9 @@ class ExponentialMovingAverage:
             decay = min(decay, (1 + self.num_updates) / (10 + self.num_updates))
         one_minus_decay = 1.0 - decay
         with torch.no_grad():
+            parameters = [p for p in parameters if p.requires_grad]
             for s_param, param in zip(self.shadow_params, parameters):
-                if param.requires_grad:
-                    s_param.sub_(one_minus_decay * (s_param - param))
+                s_param.sub_(one_minus_decay * (s_param - param))
 
     def copy_to(self, parameters):
         """