Skip to content

Commit 523d1e7

Browse files
committed
perf(model): 保存cpu权重
1 parent acad84c commit 523d1e7

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

py/lib/utils/file.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ def parse_location_xml(xml_path):
7373

7474

7575
def save_model(model, model_save_path):
76-
torch.save(model.state_dict(), model_save_path)
76+
torch.save(model.state_dict().cpu(), model_save_path)
7777

7878

7979
def save_checkpoint(model_save_path, epoch, model, optimizer, loss):

0 commit comments

Comments
 (0)