Spaces:

leuschnm
/

CrowdCounting-with-Scale-Adaptive-Selection-SASNet

Sleeping

App Files Files Community

leuschnm commited on Jan 23, 2023

Commit

9254d1a

·

1 Parent(s): 18a2e37

add cpu load

Files changed (1) hide show

app.py +1 -1

app.py CHANGED Viewed

@@ -67,7 +67,7 @@ def predict(img):
     #model = SASNet().cuda()
     model_path = "./SHHA.pth"
     # load the trained model
-    model.load_state_dict(torch.load(model_path))
     print('successfully load model from', model_path)
     with torch.no_grad():

     #model = SASNet().cuda()
     model_path = "./SHHA.pth"
     # load the trained model
+    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
     print('successfully load model from', model_path)
     with torch.no_grad():