Commit 3d287dd3 authored by zhangjx1991's avatar zhangjx1991
Browse files

bugfix: 单卡情况下,model变量没有module属性

parent cabdeed9
......@@ -203,7 +203,14 @@ if __name__=='__main__':
now = datetime.now()
timestamp = now.strftime("%Y%m%d_%H%M%S")
new_model_path = f'./saves/umt5_{timestamp}'
model.module.save_pretrained(new_model_path)
if hasattr(model,'module'):
# 多卡
model.module.save_pretrained(new_model_path)
else:
# 单卡
model.save_pretrained(new_model_path)
tokenizer.save_pretrained(new_model_path)
print("Done!")
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment