This commit is contained in:
lyuxiang.lx
2025-02-06 16:07:13 +08:00
parent 24f796a2b1
commit 2a3e033ee1
17 changed files with 187 additions and 135 deletions

View File

@@ -75,7 +75,7 @@ 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():
if k not in avg.keys() and k not in ['step', 'epoch']:
avg[k] = states[k].clone()
else:
avg[k] += states[k]