soumyamindfire

> If you make the following modifications to the function `build_res_unet`, it should work properly. > > ``` > def build_res_unet(n_input=1, n_output=2, size=256): > device = torch.device("cuda" if torch.cuda.is_available() else...

> Hi Moein, It looks like there is some difference in keys in the saved model and model architecture. While using the infer file, I am getting the following error....

Initialize the ResNet weights as well. You have to train and save the weights by yourself.