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

api_server.py: add missing device

Browse Source
This commit is contained in:
Nils Koch 2022-05-31 11:56:54 +02:00
parent 8f0c3a32b8
commit ed71c16912
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
api_server.py
View File

@ -37,7 +37,7 @@ model_path = "../train_log/models/latest.pth"
device = torch.device('cuda:0')
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)
state_dict = torch.load(model_path)['state_dict']
model.load_state_dict(state_dict, strict=True)