**I am using MSE loss() for a regression problem but with images. I want to be able to predict the outcome from a test set of images using my saved model. (Given below is a sample code i used but didnt work well). Also i am unsure if it is the right method to go about it.
How can i do that? (Sample code would be helpful)
During training my my training and validation loss doesnt improve much. Thanks**
#load saved model
device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)
model = UNet1()
checkpoint=torch.load(‘DB_model.pth’)
model.load_state_dict(checkpoint[‘state_dict’])
model.eval()
#load image and mask
image = cv2.imread("/content/image0/img2.png",0)
image = np.expand_dims(image, axis=0)
image = np.expand_dims(image, axis=1)
mask = cv2.imread("/content/wrapped/img2.png",0)
images = torch.from_numpy(image).float()
#Predictions
_mask = model(images)
#_, _mask = torch.max(_mask, dim=1)
print(_mask.shape )
fig = plt.figure(figsize=(20, 12))

setting values to rows and column variables
rows = 1
columns = 3

Adds a subplot at the 1st position
fig.add_subplot(rows, columns, 1)
plt.imshow(images.squeeze(0).permute(1,2,0).squeeze(2).numpy(), cmap=‘gray’)
plt.axis(‘off’)
plt.title(“Image”)

Adds a subplot at the 2nd position
fig.add_subplot(rows, columns, 2)
plt.imshow(_mask.squeeze(0).squeeze(0).detach().cpu().numpy(), cmap=‘gray’)
plt.axis(‘off’)
plt.title(“Prediction”)

Adds a subplot at the 3rd position
fig.add_subplot(rows, columns, 3)
plt.imshow(mask, cmap=‘gray’)
plt.axis(‘off’)
plt.title(“Mask”)