Update app.py

app.py

@@ -15,10 +15,9 @@ def predict(input_img):
         drop_rate=0.5,
         num_classes=1,)
     
-    loc = 'cuda:{}'.format(0)
-    checkpoint1 = torch.load("./machine_full_best.tar", map_location=loc)
+    checkpoint1 = torch.load("./machine_full_best.tar")
     model1.load_state_dict(checkpoint1['state_dict'])
-    checkpoint2 = torch.load("./human_full_best.tar", map_location=loc)
+    checkpoint2 = torch.load("./human_full_best.tar")
     model2.load_state_dict(checkpoint2['state_dict'])
     
     my_transform = transforms.Compose([