Hey @Usama_Hasan,
I still have a question though; the image I have is a float64 with size (256,256). When I plot it it looks fine.
When loading the model I am doing this:
forward_model.load_state_dict(torch.load(saved_models_path + '\\name_model.pt'))
forward_model.eval()
output = forward_model(image)

How should I define and reshape the image then since the image is already at the right size of 256x256?