diff --git a/cosyvoice/bin/average_model.py b/cosyvoice/bin/average_model.py
index 8843be8..b7140c1 100644
--- a/cosyvoice/bin/average_model.py
+++ b/cosyvoice/bin/average_model.py
@@ -75,10 +75,11 @@ def main():
         print('Processing {}'.format(path))
         states = torch.load(path, map_location=torch.device('cpu'))
         for k in states.keys():
-            if k not in avg.keys() and k not in ['step', 'epoch']:
-                avg[k] = states[k].clone()
-            else:
-                avg[k] += states[k]
+            if k not in ['step', 'epoch']:
+                if k not in avg.keys():
+                    avg[k] = states[k].clone()
+                else:
+                    avg[k] += states[k]
     # average
     for k in avg.keys():
         if avg[k] is not None: