Wrong predictions and model.eval() causes problem for some reason

[Himanshunitrr]

I might be doing something wrong, so wanted to know what can I correct. I am using Semantic segmentation model to do segmentation but it gives various segmentations on white patches (or non tissue patches).

without model.eval()

pretrained_weights = "/data/hmaurya/hmaurya/fcn_resnet50_unet-bcss.pth"
model = UNetModel(num_input_channels=3, num_output_channels=5, decoder_block=(3, 3))
saved_state_dict = torch.load(pretrained_weights, map_location="cpu")

print(saved_state_dict.keys())
model.load_state_dict(saved_state_dict, strict=False)
image = Image.open("/data/hmaurya/hmaurya/fibi/tilak_banks/h_e_bank/background/style_background_bank.png")
image = image.resize((2048, 2048))
image = np.array(image)[:,:,:3]
image_c = image.copy()
image = image/255.0
image = torch.tensor(image)
image = image.unsqueeze(0)
model = model.to("cuda:0")
temp_output = UNetModel.infer_batch(model, image, device="cuda:0")
temp_output = temp_output[0]

temp_output = temp_output.squeeze(0)
temp_output = np.argmax(temp_output, axis=-1)
print(np.unique(temp_output))
tumor = (temp_output == 0).astype(np.uint8)*255
stroma = (temp_output == 1).astype(np.uint8)*255
inflam = (temp_output == 2).astype(np.uint8)*255
necorsis = (temp_output == 3).astype(np.uint8)*255
others = (temp_output == 4).astype(np.uint8)*255
fig, ax = plt.subplots(1, 6, figsize=(20, 10))
ax[0].imshow(image_c)
ax[1].imshow(tumor)
ax[2].imshow(stroma)
ax[3].imshow(inflam)
ax[4].imshow(necorsis)
ax[5].imshow(others)

ax[0].set_title("Original Image")
ax[1].set_title("Tumor")
ax[2].set_title("Stroma")
ax[3].set_title("Inflammation")
ax[4].set_title("Necrosis")
ax[5].set_title("Others")

plt.show()

with mode.eval():