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
Show 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,6 +15,7 @@ class ExponentialMovingAverageModel:
        dtype_grouped_names = dict()
        ordered_dtype = []
        for n, p in model.named_parameters():
+            if p.requires_grad:
                if p.dtype not in dtype_grouped_names:
                    dtype_grouped_names[p.dtype] = []
                    ordered_dtype.append(p.dtype)