map to cpu

CatVsDogTrain.py

@@ -224,7 +224,7 @@ def train_and_validate(model, loss_criterion, optimizer, epochs=25):
         )
 
         # Save if the model has best accuracy till now
-        torch.save(model, "TrainLoopImproveCatsDogs.pth")
+        torch.save(model, "CatVsDogsModel.pth")
 
     return model, history
 
@@ -258,7 +258,7 @@ history = []
 def test():
     test = datasets.ImageFolder(root="PetTest/", transform=convert)
     testLoader = DataLoader(test, batch_size=16, shuffle=False)
-    checkpoint = torch.load("catsvdogs.pth")
+    checkpoint = torch.load("CatVsDogsModel.pth")
     alexnet.load_state_dict(checkpoint["state_dict"])
     optimizer.load_state_dict(checkpoint["optimizer"])
     for params in alexnet.parameters():
@@ -388,7 +388,7 @@ def predict(model, test_image_name):
 
 if PREDICT:
     checkpoint = torch.load(
-        "ImprovedCatVsDogsModel.pth", map_location=torch.device("cpu")
+        "CatVsDogsModel.pth", map_location=torch.device("cpu")
     )
     alexnet.load_state_dict(checkpoint["state_dict"])
     alexnet = alexnet.to(device)

app.py

@@ -5,7 +5,7 @@ from timm.data import resolve_data_config
 from timm.data.transforms_factory import create_transform 
 
 LABELS = {0:'Cat', 1:'Dog'}
-model = torch.load('CatVsDogsModel.pth')
+model = torch.load('CatVsDogsModel.pth',map_location='cpu')
 model.eval()
 transform = create_transform(**resolve_data_config({},model=model))