torch多卡模型转单卡

from collections import OrderedDict
state_dict1 = torch.load(self.model1)
module_state_dict1 = OrderedDict()
for k, v in state_dict1.items():
	name = k[7:]
	module_state_dict1[name] = v
	self.model = self.model.cuda()        	
self.model.load_state_dict(module_state_dict1, strict=True)
#self.model = self.model.cuda()
self.model.eval()

猜你喜欢

转载自blog.csdn.net/caihuanqia/article/details/113711554