Spaces:

chloecodes
/

sam

Build error

Nguyen Thai Thao Uyen commited on Apr 28, 2024

Commit

a541988

1 Parent(s): c3f5b96

run.py Fix cache permission issue

Files changed (1) hide show

run.py CHANGED Viewed

@@ -16,7 +16,7 @@ def pred(src):
     model_config = SamConfig.from_pretrained("facebook/sam-vit-base", cache_dir=cache_dir)
     model = SamModel(config=model_config)
-    model.load_state_dict(torch.load(MODEL_FILE))
     with open("sam-config.json", "r") as f: # modified config json file
         modified_config_dict = json.load(f)

     model_config = SamConfig.from_pretrained("facebook/sam-vit-base", cache_dir=cache_dir)
     model = SamModel(config=model_config)
+    model.load_state_dict(torch.load(MODEL_FILE, map_location=torch.device('cpu')))
     with open("sam-config.json", "r") as f: # modified config json file
         modified_config_dict = json.load(f)