Amrit_Das:

num_ftrs = model_ft.fc.in_features model_ft.fc = nn.Linear(num_ftrs, 4)


Here, you have updated the fc layer on resnet18.
Your saved model and loading models are different.
Code for prediction should be as follows:
checkpoint = torch.load("./models/custom_model13.model")
# Load model here
model = resnet18(pretrained=True)
# make the fc layer similar to the saved model
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 4)
# Now load the checkpoint
model.load_state_dict(checkpoint)
model.eval()