Below is how I loaded my pre-trained model and created test dataloader.
from collections import Counter, OrderedDict

x = torch.load("pretrained_att_resnet50.pth")

torch.manual_seed(42)
torch.cuda.manual_seed(42)
random.seed(42)
np.random.seed(42)

model = TwoAttresnet50()
model.fc = nn.Linear(512*4,2)
model.load_state_dict(x)
model = model.to(device)

class CustomDataset_test(Dataset):
    def __init__(self, files, labels, augmentation, transforms2, valid):
        self.files = files
        self.labels = labels
        self.aug = augmentation
        self.transforms2 = transforms2
        self.valid = valid
        self.data = []
        
        for i in range(len(self.files)):
            sample = {}
            sample['img'] = Image.open(self.files[i])
            sample['label'] = labels[i]
            sample['fname'] = self.files[i]
            self.data.append(sample)
        
    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        x = self.data[idx]['img']
        y = self.data[idx]['label']
        
        x1 = self.aug(x)
        x2 = self.transforms2(x)
        return {"img": np.array(x1, dtype='float32'), "labels": y, "original": np.array(x2, dtype='float32'),"fname": self.data[idx]['fname']}
    
test_transforms = transforms.Compose([
                    #transforms.Resize(224),
                    transforms.ToTensor(),
                    transforms.Normalize((0.2880, 0.2441, 0.2705),(0.2033, 0.1889, 0.1822))
                    #transforms.Normalize((0.2206, 0.1899, 0.2003), (0.1690, 0.1592, 0.1524))
                    ])

original_transforms = transforms.Compose([
    #transforms.Resize(224)
])

torch.manual_seed(42)
torch.cuda.manual_seed(42)
random.seed(42)
np.random.seed(42)

batch_size_test = 1
test_dataset = CustomDataset_test(test_files,test_labels,augmentation=test_transforms, transforms2 = original_transforms, valid=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size_test, shuffle=False, drop_last = True, num_workers =1)

print('######### Dataset class created #########')
print('Number of images: ', len(test_files))
print("test size in no of batch: ",len(test_loader))
print("test size: ",len(test_loader)*batch_size_test)