m

handler.py

@@ -1,11 +1,10 @@
 from typing import Dict, List, Any
-from comet import load_from_checkpoint as lfc
-import os
+from comet import load_from_checkpoint
 
 
 class EndpointHandler():
     def __init__(self, path=""):
-        self.m = lfc("model.ckpt")
+        self.model = load_from_checkpoint("model.ckpt")
 
     def __call__(self, data: Dict[str, Any]) -> List[Any]:
         """
@@ -22,7 +21,7 @@ class EndpointHandler():
 
         data = inputs.pop("data")
         
-        model_output = self.m.predict(data, batch_size=batch_size, num_workers=workers, gpus=0)
+        model_output = self.model.predict(data, batch_size=batch_size, num_workers=workers, gpus=0)
         scores = model_output["scores"]
 
         return scores