Yiheng Hu commited on
Commit
ced1bec
·
1 Parent(s): c9d8c8e

update handler

Browse files
Files changed (1) hide show
  1. handler.py +2 -2
handler.py CHANGED
@@ -8,7 +8,7 @@ class EndpointHandler():
8
  # Preload all the elements you are going to need at inference.
9
  # pseudo:
10
  # self.model= load_model(path)
11
- self.md_model = AltCLIPModel.from_pretrained(path)
12
  self.md_processor = AltCLIPProcessor.from_pretrained(path)
13
 
14
 
@@ -25,6 +25,6 @@ class EndpointHandler():
25
  # self.model(input)
26
  with torch.inference_mode():
27
  texts = data.pop("inputs",data)
28
- inputs = self.md_processor(text = texts, padding=True, return_tensors="pt")
29
  text_feature = self.md_model.get_text_features(**inputs)
30
  return {"feature":text_feature}
 
8
  # Preload all the elements you are going to need at inference.
9
  # pseudo:
10
  # self.model= load_model(path)
11
+ self.md_model = AltCLIPModel.from_pretrained(path).to(device)
12
  self.md_processor = AltCLIPProcessor.from_pretrained(path)
13
 
14
 
 
25
  # self.model(input)
26
  with torch.inference_mode():
27
  texts = data.pop("inputs",data)
28
+ inputs = self.md_processor(text = texts, padding=True, return_tensors="pt").to(device)
29
  text_feature = self.md_model.get_text_features(**inputs)
30
  return {"feature":text_feature}