Log device
Browse files- 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 |
"""
|