I have trained my UNet model and got a good accuracy result though for a single image prediction I am suffering more and more. I am using google colab for training purposes and saving my model into google drive. After loading the trained model prediction is not good, say only about 30/40 % though I was getting 90% accuracy during training. What might be the possible cause for this?
My training loop:
for batch in train_loader:
    model.train()
    images, labels = batch
    labels = labels.squeeze(dim=1)
    labels = labels.to(torch.long)

    iteration += 1

    preds = model(images)
    loss = criterion(preds, labels)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    predicted = torch.argmax(preds.data, 1)
    total_train += labels.nelement()
    correct_train += predicted.eq(labels.data).sum().item()
    train_accuracy = 100 * correct_train / total_train

from PIL import Image
from torchvision import transforms
input_image = Image.open('drive/My Drive/BackgroundImageRemoval/Test Images/400pixel.jpg')
preprocess = transforms.Compose([
    transforms.CenterCrop(248),
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

input_tensor = preprocess(input_image)

input_batch = input_tensor.unsqueeze(0)
with torch.no_grad():
  output2 = model(input_batch)
  output2 = output2.squeeze(0)
output2_predictions = output2.argmax(0)
r = Image.fromarray(output2_predictions.byte().cpu().numpy()).resize(input_image.size)
r