@ptrblck ,
I’ve traced the error to model.load_state_dict( ). To give you context,
import torch.nn as nn
from torchvision.models.resnet import ResNet,BasicBlock

class ImageClassifier(ResNet):
	def __init__(self):
		super(ImageClassifier,self).__init__(BasicBlock,[2,2,2,2],num_classes=2)
		self.fc = nn.Linear(512,2)

import torch
model = ImageClassifier()
model.load_state_dict(torch.load("trainedRes18.pth"))

I’ve trained and saved the “trainedRes18.pth” using:
The code snippet I’ve posted.
torch.save(model_ft.state_dict(),"trainedRes18.pth")