Guill-Bla commited on
Commit
3741946
·
verified ·
1 Parent(s): 88bb0df

Update tasks/image.py

Browse files
Files changed (1) hide show
  1. tasks/image.py +1 -1
tasks/image.py CHANGED
@@ -25,7 +25,7 @@ ROUTE = "/image"
25
 
26
  device = "cpu"
27
  model_path = "mobilevit_segmentation_full_data.pth"
28
- feature_extractor = MobileViTImageProcessor.from_pretrained("apple/deeplabv3-mobilevit-xx-small").to(device)
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()
 
25
 
26
  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()