Update handler.py
Browse files- handler.py +2 -0
handler.py
CHANGED
|
@@ -28,6 +28,8 @@ class EndpointHandler:
|
|
| 28 |
)
|
| 29 |
|
| 30 |
self.tokenizer = AutoTokenizer.from_pretrained(path)
|
|
|
|
|
|
|
| 31 |
|
| 32 |
def __call__(self, data: Any) -> List[List[Dict[str, str]]]:
|
| 33 |
|
|
|
|
| 28 |
)
|
| 29 |
|
| 30 |
self.tokenizer = AutoTokenizer.from_pretrained(path)
|
| 31 |
+
print("Memory footprint: ", self.model.get_memory_footprint())
|
| 32 |
+
print("Device map: ", model.hf_device_map)
|
| 33 |
|
| 34 |
def __call__(self, data: Any) -> List[List[Dict[str, str]]]:
|
| 35 |
|