map location model to cpu

test_few_shot.py

@@ -31,7 +31,7 @@ def test_main_model(opts):
         st.write("Loading Model Weight...")
     model_main = ModelMain(opts)
     path_ckpt = os.path.join(f"{opts.model_path}")
-    model_main.load_state_dict(torch.load(path_ckpt)['model'], map_location=torch.device('cpu'))
+    model_main.load_state_dict(torch.load(path_ckpt, map_location=torch.device('cpu'))['model'])
     model_main.to(device)
     model_main.eval()
     with torch.no_grad():