I tried to load pretrained MNIST model on a CPU only machine, however, I met
RuntimeError: Attempting to deserialize object on a CUDA device but
torch.cuda.is_available() is False. If you are running on a CPU-only machine,
please use torch.load with map_location='cpu' to map your storages to the CPU.
So I added a parameter cuda with a default value True when creating models. This parameter is only used when loading pretrained models from model_zoo.
When load model from model_zoo.load_url, I set the map_location to "cuda" if cuda=True in the parameter, else "cpu"
map_location (optional) – a function or a dict specifying how to remap storage locations (see torch.load)
I tried to load pretrained MNIST model on a CPU only machine, however, I met
So I added a parameter
cudawith a default valueTruewhen creating models. This parameter is only used when loading pretrained models from model_zoo.When load model from model_zoo.load_url, I set the
map_locationto "cuda" if cuda=True in the parameter, else "cpu"from https://pytorch.org/docs/stable/model_zoo.html