def load_model(path):
    """Loads model and return it without DataParallel table."""
    if os.path.isfile(path):
        print("=> loading checkpoint '{}'".format(path))
        checkpoint = torch.load(path)

        for key in checkpoint["state_dict"]:
            print(key)

        # size of the top layer
        N = checkpoint["state_dict"]["decoder.out_conv.bias"].size()

        # build skeleton of the model
        sob = "sobel.0.weight" in checkpoint["state_dict"].keys()
        model = models.__dict__[checkpoint["arch"]](sobel=sob, out=int(N[0]))

        # deal with a dataparallel table
        def rename_key(key):
            if not "module" in key:
                return key
            return "".join(key.split(".module"))

        checkpoint["state_dict"] = {
            rename_key(key): val for key, val in checkpoint["state_dict"].items()
        }

        # load weights
        model.load_state_dict(checkpoint["state_dict"])
        print("Loaded")
    else:
        model = None
        print("=> no checkpoint found at '{}'".format(path))
    return model
Loads model and return it without DataParallel table.
code/utils/util.py