submission-template-mobilevit

Sleeping

Guill-Bla commited on Jan 29

Commit

cd9b6a5

verified ·

1 Parent(s): 3741946

Update tasks/image.py

Files changed (1) hide show

tasks/image.py CHANGED Viewed

@@ -27,7 +27,7 @@ device = "cpu"
 model_path = "mobilevit_segmentation_full_data.pth"
 feature_extractor = MobileViTImageProcessor.from_pretrained("apple/deeplabv3-mobilevit-xx-small")
 model = MobileViTForSemanticSegmentation.from_pretrained("apple/deeplabv3-mobilevit-xx-small").to(device)
-model.load_state_dict(torch.load(model_path), map_location=torch.device('cpu'))
 model.eval()
 def preprocess(image):

 model_path = "mobilevit_segmentation_full_data.pth"
 feature_extractor = MobileViTImageProcessor.from_pretrained("apple/deeplabv3-mobilevit-xx-small")
 model = MobileViTForSemanticSegmentation.from_pretrained("apple/deeplabv3-mobilevit-xx-small").to(device)
+model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
 model.eval()
 def preprocess(image):