Hugging Face's logo

Spaces:
frugal-ai-challenge
/
image-Guill-Bla
Sleeping

App Files Community
Guill-Bla commited on
Commit
cd9b6a5
·
verified ·
1 Parent(s): 3741946

Update tasks/image.py

Browse files
Files changed (1) hide show
  1. tasks/image.py +1 -1
tasks/image.py CHANGED
@@ -27,7 +27,7 @@ device = "cpu"
27
  model_path = "mobilevit_segmentation_full_data.pth"
28
  feature_extractor = MobileViTImageProcessor.from_pretrained("apple/deeplabv3-mobilevit-xx-small")
29
  model = MobileViTForSemanticSegmentation.from_pretrained("apple/deeplabv3-mobilevit-xx-small").to(device)
30
- model.load_state_dict(torch.load(model_path), map_location=torch.device('cpu'))
31
  model.eval()
32
 
33
  def preprocess(image):
 
27
  model_path = "mobilevit_segmentation_full_data.pth"
28
  feature_extractor = MobileViTImageProcessor.from_pretrained("apple/deeplabv3-mobilevit-xx-small")
29
  model = MobileViTForSemanticSegmentation.from_pretrained("apple/deeplabv3-mobilevit-xx-small").to(device)
30
+ model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
31
  model.eval()
32
 
33
  def preprocess(image):