RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.

[slivara]

(base) C:\Users\B4-410\lpthw\LearningToCompare_FSL-master\miniimagenet>python miniimagenet_train_one_shot.py -w 5 -s 1 -b 15
init data folders
init neural networks
Traceback (most recent call last):
File "miniimagenet_train_one_shot.py", line 269, in
main()
File "miniimagenet_train_one_shot.py", line 150, in main
feature_encoder.load_state_dict(torch.load(str("./models/miniimagenet_feature_encoder_" + str(CLASS_NUM) +"way_" + str(SAMPLE_NUM_PER_CLASS) +"shot.pkl")))
File "D:\Anaconda3\lib\site-packages\torch\serialization.py", line 426, in load
return _load(f, map_location, pickle_module, **pickle_load_args)
File "D:\Anaconda3\lib\site-packages\torch\serialization.py", line 613, in _load
result = unpickler.load()
File "D:\Anaconda3\lib\site-packages\torch\serialization.py", line 576, in persistent_load
deserialized_objects[root_key] = restore_location(obj, location)
File "D:\Anaconda3\lib\site-packages\torch\serialization.py", line 155, in default_restore_location
result = fn(storage, location)
File "D:\Anaconda3\lib\site-packages\torch\serialization.py", line 131, in _cuda_deserialize
device = validate_cuda_device(location)
File "D:\Anaconda3\lib\site-packages\torch\serialization.py", line 125, in validate_cuda_device
device, torch.cuda.device_count()))
RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count() is 1. Please use torch.load with map_location to map your storages to an existing device.

[githubwys]

i met the same problem.

you have to assign a certain gpu.

find the file miniimagenet_train_one_shot.py
from :
#feature_encoder.load_state_dict(torch.load(str("./models/miniimagenet_feature_encoder_" + str(CLASS_NUM) +"way_" + str(SAMPLE_NUM_PER_CLASS) +"shot.pkl")))
to :
feature_encoder.load_state_dict(torch.load(str("./models/miniimagenet_feature_encoder_" + str(CLASS_NUM) +"way_" + str(SAMPLE_NUM_PER_CLASS) +"shot.pkl"), map_location='cuda:0'))

also the same as the relation_network part.

hope helpful to you