fix bug

FunAudioLLM · Feb 11, 2025 · 8903005 · 8903005
1 parent f77c6a8
commit 8903005
Showing 1 changed file with 5 additions and 4 deletions.
diff --git 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: