foamliu/Deep-Image-Matting-PyTorch

Run trained model on CPU

patriotyk opened this issue · 1 comments

I try to run eval.py on my laptop that doesn't have cuda and I got the following error:

Traceback (most recent call last):
  File "eval.py", line 62, in <module>
    pred = model(x)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 547, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/torch/nn/parallel/data_parallel.py", line 146, in forward
    "them on device: {}".format(self.src_device_obj, t.device))
RuntimeError: module must have its parameters and buffers on device cuda:0 (device_ids[0]) but found one of them on device: cpu

I also modified the load state code to
checkpoint = torch.load(checkpoint, map_location=lambda storage, loc: storage)
but it doesn't fix all issues.

Fixed, feel free to activate this if it repro.