Spaces:

yijiu
/

pose_experiment

Running

yijiu commited on Dec 16, 2023

Commit

44207e0

1 Parent(s): 6a5f391

fix:load model in cpu

Files changed (1) hide show

app.py CHANGED Viewed

@@ -21,7 +21,7 @@ def predict(numpy_img):
     #choose model class hr_w32
     model=hr_w32().to(device)
     #load weights of model
-    model.load_state_dict(torch.load('./weights/HRNet_epoch20_loss0.000474.pth')['model'])
     # #set model to pred state
     model.eval()
     # #predict the heatmaps of joints

     #choose model class hr_w32
     model=hr_w32().to(device)
     #load weights of model
+    model.load_state_dict(torch.load('./weights/HRNet_epoch20_loss0.000474.pth',map_location=torch.device('cpu'))['model'])
     # #set model to pred state
     model.eval()
     # #predict the heatmaps of joints