fixed map location issue in models.py

models.py

@@ -3,13 +3,13 @@ from torchvision.models import  ViT_B_16_Weights, EfficientNet_B2_Weights
 
 
 def get_vit_16_base_transformer():
-    vit_b_16_model = torch.load(r"models\ViT_16_base_101_classes_pretrained_custom_head.pth")
+    vit_b_16_model = torch.load("models/ViT_16_base_101_classes_pretrained_custom_head.pth", map_location = torch.device("cpu"))
     vit_b_16_transforms = ViT_B_16_Weights.DEFAULT.transforms()
 
     return vit_b_16_model, vit_b_16_transforms
 
 def get_effnet_b2():
-    eff_net_b2_model = torch.load(r"models\effnet_b2_101_classes_pretrained_custom_head.pth")
+    eff_net_b2_model = torch.load("models/effnet_b2_101_classes_pretrained_custom_head.pth", map_location = torch.device("cpu"))
     eff_net_b2_transforms = EfficientNet_B2_Weights.DEFAULT.transforms()
 
     return eff_net_b2_model, eff_net_b2_transforms