Cpt.Captain
/
CREStereo-pytorch-nxt
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

api_server.py: add missing device

Browse Source
main
Nils Koch 2 years ago
parent 8f0c3a32b8
commit ed71c16912
1 changed files with 1 additions and 1 deletions
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Show Stats Download Patch File Download Diff File
  1. 2
      api_server.py

2
api_server.py
Unescape View File

@ -37,7 +37,7 @@ model_path = "../train_log/models/latest.pth"
device = torch.device('cuda:0') device = torch.device('cuda:0')
model = Model(max_disp=256, mixed_precision=False, test_mode=True) model = Model(max_disp=256, mixed_precision=False, test_mode=True)
model = nn.DataParallel(model, device_ids=[]) model = nn.DataParallel(model, device_ids=[device])
# model.load_state_dict(torch.load(model_path), strict=False) # model.load_state_dict(torch.load(model_path), strict=False)
state_dict = torch.load(model_path)['state_dict'] state_dict = torch.load(model_path)['state_dict']
model.load_state_dict(state_dict, strict=True) model.load_state_dict(state_dict, strict=True)

Write Preview
Loading…
Cancel
Save