Spaces:

CUHKWilliam
/

Multiple-shots

Sleeping

CUHKWilliam commited on Jun 4, 2024

Commit

6bd32e7

verified ·

1 Parent(s): 5f5377c

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -24,7 +24,7 @@ def inference_mask1(
     model = DCAMA('resnet50', 'resnet50_a1h-35c100f8.pth', True)
     model.eval()
     model.cpu()
-    state_dict = torch.load('model_45.pt')
     if 'state_dict' in state_dict.keys():
         state_dict = state_dict['state_dict']
     state_dict2 = {}

     model = DCAMA('resnet50', 'resnet50_a1h-35c100f8.pth', True)
     model.eval()
     model.cpu()
+    state_dict = torch.load('model_45.pt', map_location=torch.device('cpu'))
     if 'state_dict' in state_dict.keys():
         state_dict = state_dict['state_dict']
     state_dict2 = {}