28914875e1a85a1c860ece7ddce715adbcd5c067,data/colorize_image.py,ColorizeImageTorch,prep_net,#ColorizeImageTorch#,215

Before Change

    def prep_net(self, gpu_id=None, path="", dist=False):
        import torch
        import models.pytorch.model as model
        device = torch.device("cuda:{}".format(gpu_id)) if gpu_id != -1 else torch.device("cpu")
        print("path = %s, device = %s" % (path, device))
        print("Model set! dist mode? ", dist)
        self.net = model.SIGGRAPHGenerator(dist=dist)
        state_dict = torch.load(path, map_location=str(device))

After Change

        for key in list(state_dict.keys()):  // need to copy keys here because we mutate in loop
            self.__patch_instance_norm_state_dict(state_dict, self.net, key.split("."))
        self.net.load_state_dict(state_dict)
        if gpu_id != -1:
            self.net.cuda()
        self.net.eval()
        self.net_set = True

    def __patch_instance_norm_state_dict(self, state_dict, module, keys, i=0):

Instances