Unable to run the model on CPU

[srikrishnag]

Was trying to run the model in cpu. Below are the changes made for the same, but facing below issue in running the model, can you help to fix this?

Error generated:

Traceback (most recent call last):
File "model.py", line 129, in
main()
File "model.py", line 40, in main
model.load_state_dict(checkpoint['state_dict'])
File "E:\anaconda\lib\site-packages\torch\nn\modules\module.py", line 719, in load_state_dict
self.class.name, "\n\t".join(error_msgs)))
RuntimeError: Error(s) in loading state_dict for DataParallel:
Missing key(s) in state_dict: "module.densenet121.features.denseblock1.denselayer1.norm1.weight", "module.densenet121.features.denseblock1.denselayer1.norm1.bias", "module.densenet121.features.denseblock1.denselayer1.norm1.running_mean", "module.densenet121.features.denseblock1.denselayer1.norm1.running_var",

Model changes:

`def main():

cudnn.benchmark = False

# initialize and load the model
model = DenseNet121(N_CLASSES).to(torch.device("cpu"))
model = torch.nn.DataParallel(model)

if os.path.isfile(CKPT_PATH):
    print("=> loading checkpoint")
    checkpoint = torch.load(CKPT_PATH, map_location=lambda storage, loc: storage)
    model.load_state_dict(checkpoint['state_dict'])
    print("=> loaded checkpoint")
else:
    print("=> no checkpoint found")

normalize = transforms.Normalize([0.485, 0.456, 0.406],
                                 [0.229, 0.224, 0.225])

test_dataset = ChestXrayDataSet(data_dir=DATA_DIR,
                                image_list_file=TEST_IMAGE_LIST,
                                transform=transforms.Compose([
                                    transforms.Resize(256),
                                    transforms.TenCrop(224),
                                    transforms.Lambda
                                    (lambda crops: torch.stack([transforms.ToTensor()(crop) for crop in crops])),
                                    transforms.Lambda
                                    (lambda crops: torch.stack([normalize(crop) for crop in crops]))
                                ]))
test_loader = DataLoader(dataset=test_dataset, batch_size=BATCH_SIZE,
                         shuffle=False, num_workers=8, pin_memory=True)

# initialize the ground truth and output tensor
gt = torch.FloatTensor()
#gt = gt.cuda()
pred = torch.FloatTensor()
#pred = pred.cuda()

# switch to evaluate mode
model.eval()`