Printing device map

handler.py

@@ -20,6 +20,9 @@ class EndpointHandler():
             device_map = infer_auto_device_map(model, no_split_module_classes=["T5Block"])
         device_map['language_model.lm_head'] = device_map["language_model.encoder.embed_tokens"]
 
+        print(device_map)
+        exit()
+        
         self.model = Blip2ForConditionalGeneration.from_pretrained(
             "Salesforce/blip2-flan-t5-xxl", device_map=device_map,
             # torch_dtype=torch.float16