search

mihaidusmanu / d2-net

D2-Net: A Trainable CNN for Joint Description and Detection of Local Features
Other
764 stars 163 forks source link

Small bug when running on CPU-only machines #51

Closed uzpaka closed 4 years ago

uzpaka commented 4 years ago

https://github.com/mihaidusmanu/d2-net/blob/81983664d1a29fabf09636b3735c5df9151db600/lib/model_test.py#L61

With this line, pytorch tries to load the weights on the GPU and throws an error if it can't. The following update should solve the issue:

if model_file is not None: if use_cuda: self.load_state_dict(torch.load(model_file)['model']) else: self.load_state_dict(torch.load(model_file, map_location='cpu')['model'])

mihaidusmanu commented 4 years ago

This should be fixed by the latest commit. Thank you!