How to test in CPU mode? about wsr

[supercaoO]

@scimg Hello! If your machine has no GPUs or your PyTorch only supports CPU mode (i.e., torch.cuda.is_available() == False), you can try the following code modification to test in CPU mode temporarily .

Specifically, modify the codes between Line 331 and Line 334 in ./solvers/SRSolver.py

checkpoint = torch.load(model_path)
if 'state_dict' in checkpoint.keys(): checkpoint = checkpoint['state_dict']
load_func = self.model.load_state_dict if isinstance(self.model, nn.DataParallel) \
    else self.model.module.load_state_dict

to

if self.use_gpu:
    checkpoint = torch.load(model_path)
    if 'state_dict' in checkpoint.keys(): checkpoint = checkpoint['state_dict']
    load_func = self.model.load_state_dict if isinstance(self.model, nn.DataParallel) \
    else self.model.module.load_state_dict
else:
    checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage)
    if 'state_dict' in checkpoint.keys(): checkpoint = checkpoint['state_dict']
    checkpoint_cpu = {}
    for key in checkpoint.keys():
        checkpoint_cpu[key.replace('module.', '')] = checkpoint[key]
    checkpoint = checkpoint_cpu
    load_func = self.model.load_state_dict

Then, try to test the pre-trained model again:

cd WSR
python test.py -opt options/test/test_WSR_Set5.json

I will update the repository to support CPU mode as soon as possible.

[scimg]

@supercaoO Thanks for replying, it works fine now.