加载预训练模型默认在第0块gpu上,将其修改为cpu的预训练加载方式:
model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')), strict=False)
加载预训练模型默认在第0块gpu上,将其修改为cpu的预训练加载方式:
model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')), strict=False)