Update handler.py

handler.py

@@ -2,12 +2,14 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import subprocess
 
+# Install flash-attn
 subprocess.run('pip install flash-attn', shell=True)
 
 class CustomModelHandler:
     def __init__(self, model_name_or_path: str):
         self.model_name_or_path = model_name_or_path
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        print(f"Using device: {self.device}")
         self.load_model()
 
     def load_model(self):
@@ -21,6 +23,16 @@ class CustomModelHandler:
             )
             self.model.to(self.device)
             print(f"Model loaded and moved to {self.device}")
+
+            # Check if the model parameters are on the GPU
+            all_on_gpu = all(param.device.type == 'cuda' for param in self.model.parameters())
+            if not all_on_gpu:
+                print("Warning: Not all model parameters are on the GPU!")
+            else:
+                print("All model parameters are on the GPU.")
+
+            # Confirm model device
+            print(f"Model is on device: {self.model.device}")
         except Exception as e:
             print(f"An error occurred while loading the model: {e}")
             raise
@@ -41,3 +53,9 @@ handler = CustomModelHandler("microsoft/Phi-3-vision-128k-instruct")
 # Example prediction function
 def predict(input_text):
     return handler.predict(input_text)
+
+# Example usage
+if __name__ == "__main__":
+    input_text = "Hello, how are you?"
+    predictions = predict(input_text)
+    print("Predictions:", predictions)