OK, I think I’ve got where the problem rises: the model weight saved with
torch.save(model.state_dict(), file)
contains device info and
torch.load(model_file) will load the weight directly into the device according to the saved device info rather than load into CPU. So, if the previously used device is short of memory, this loading process will crash.