YaYaB
/

yb_test_inference_endpoint_image_cls

Model card Files Files and versions

YassineB commited on Oct 17, 2022

Commit

6207ffd

·

1 Parent(s): c3ac0f4

fix output

Files changed (1) hide show

handler.py +2 -1

handler.py CHANGED Viewed

@@ -30,4 +30,5 @@ class EndpointHandler():
         batch = self.preprocess(img).unsqueeze(0)
         prediction = self.pipeline(batch).squeeze(0).softmax(0)
-        return prediction

         batch = self.preprocess(img).unsqueeze(0)
         prediction = self.pipeline(batch).squeeze(0).softmax(0)
+        return prediction.tolist()