nnunetv2，save weight and network to a pth file

[qwertyuiop194]

unet = nnUNetTrainer(plans=plan_json, configuration='3d_fullres', fold=0, dataset_json=data_json)
unet.initialize()
#print(unet.network)
#print(model_weight)
unet.network.load_state_dict(model_weight,strict=False)
trace_model = torch.jit.trace(unet.network, torch.randn(pp_patchsize).to("cuda"))
torch.jit.save(trace_model, os.path.join(output, 'nnunet.jit.pt'))

Execute the above code, unable to save, please help，thanks

[qwertyuiop194]

Excuse me again, do you have a solution? Thank you very much

[anxingle]

Can you give more detail about your issue？ like wrong info