filter params in ema (#19)

f2b0b488 · Jinhua Zhu · GitHub · 4e4dc4b8 · f2b0b488
Unverified Commit f2b0b488 authored Feb 23, 2023 by Jinhua Zhu Committed by GitHub Feb 23, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 4 deletions

unicore/ema.py unicore/ema.py +5 -4

No files found.
--- a/unicore/ema.py
+++ b/unicore/ema.py
@@ -15,10 +15,11 @@ class ExponentialMovingAverageModel:
        dtype_grouped_names = dict()
        ordered_dtype = []
        for n, p in model.named_parameters():
-            if p.dtype not in dtype_grouped_names:
-                dtype_grouped_names[p.dtype] = []
-                ordered_dtype.append(p.dtype)
-            dtype_grouped_names[p.dtype].append(n)
+            if p.requires_grad:
+                if p.dtype not in dtype_grouped_names:
+                    dtype_grouped_names[p.dtype] = []
+                    ordered_dtype.append(p.dtype)
+                dtype_grouped_names[p.dtype].append(n)

        ordered_names = list(chain(*(dtype_grouped_names[n] for n in ordered_dtype)))