Hello together,
I have successfully converted my model to “script”. Only when I test my model now, the weights do not match, so my program does not work. Does anyone have any idea?
import argparse

import os
from os import path
from unifuse import UniFuse
import torch


def load_model(model, args):
    if args.v:
        print("Load Model")
    model_dict = model.state_dict()
    pretrained_dict = torch.load(args.model_path)
    pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
    model_dict.update(pretrained_dict)
    return model

def convert(args):
    model = UniFuse(18, args.height, args.width, True, args.max_depth, "cee", True)
    model = load_model(model, args)
    if args.cuda:
        model.cuda()
        if args.v:
            print("Set model GPU")
    else:
        model.cpu()
        if args.v:
            print("Set model CPU")

    my_script_module = torch.jit.script(model)
    save(my_script_module, args)
    if args.v:
        print(my_script_module.code)

def save(script, args):
    name = "wrapped.pt"
    file = os.path.join(args.output_path, name)
    torch.jit.save(script, file)

def main():
    parser = argparse.ArgumentParser(description="Convert the trained model with torchscript into a binary format")
    # path to the data
    parser.add_argument("--model_path", default="/model.pth", type=str, help="Path to the trained model")
    parser.add_argument("--output_path", default=".", type=str, help="Output path where it will be stored")

    # model params
    parser.add_argument("--num_layers", default=18, type=int, help="Number on layers will be needed")
    parser.add_argument("--height", default=512, type=int, help="Define the height of the input and output Image")
    parser.add_argument("--width", default=1024, type=int, help="Define the width of the input and output")
    parser.add_argument("--max_depth", default=10.0, type=float, help="Define the max depth")
    parser.add_argument("--cuda", default=True, action="store_true", help="if set, use cuda")

    # show the logs
    parser.add_argument("-v", default=True, action="store_true", help="if set, shows the logs")

    args = parser.parse_args()

    if not path.exists(args.model_path):
        exit("Model path didnt exists!")
    if not path.exists(args.output_path):
        exit("Output path didnt exists!")

    convert(args)

if __name__ == "__main__":
    main()