if len(self.gpu_ids) > 0 and torch.cuda.is_available():
net.module.load_state_dict(torch.load(save_path))
else:
net.load_state_dict(torch.load(save_path))
// print network information
def print_networks(self, verbose):
print("---------- Networks initialized -------------")
After Change
if isinstance(name, str):
save_filename = "%s_net_%s.pth" % (which_epoch, name)
save_path = os.path.join(self.save_dir, save_filename)
net = getattr(self, "net" + name)
if isinstance(net, torch.nn.DataParallel):
net = net.module
// if you are using PyTorch newer than 0.4 (e.g., built from
// GitHub source), you can remove str() on self.device
state_dict = torch.load(save_path, map_location=str(self.device))
// patch InstanceNorm checkpoints prior to 0.4
for key in state_dict:
self.__patch_instance_norm_state_dict(state_dict, net, key.split("."))
net.load_state_dict(state_dict)
// print network information
def print_networks(self, verbose):
print("---------- Networks initialized -------------")
