jpohhhh commited on
Commit
dc77190
·
1 Parent(s): 47bfcfc

Log device

Browse files
Files changed (1) hide show
  1. handler.py +1 -0
handler.py CHANGED
@@ -14,6 +14,7 @@ class EndpointHandler():
14
  self.tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/msmarco-MiniLM-L-6-v3')
15
  self.model = AutoModel.from_pretrained('sentence-transformers/msmarco-MiniLM-L-6-v3')
16
  self.model.to(self.device)
 
17
 
18
  def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
19
  """
 
14
  self.tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/msmarco-MiniLM-L-6-v3')
15
  self.model = AutoModel.from_pretrained('sentence-transformers/msmarco-MiniLM-L-6-v3')
16
  self.model.to(self.device)
17
+ print("model will run on ", self.device)
18
 
19
  def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
20
  """