Update app.py

app.py

@@ -4,7 +4,14 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 device = "cuda" if torch.cuda.is_available() else "cpu"  # Automatically detect GPU or CPU
 model_name = "tanusrich/Mental_Health_Chatbot"
 
-model = AutoModelForCausalLM.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.float16,  # Reduce memory usage
+    device_map="auto",  # Automatically assigns to GPU if available
+    low_cpu_mem_usage=True,
+    max_memory={0: "4GiB", "cpu": "6GiB"}, # Optimize CPU memory
+    offload_folder=None
+)
 
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 '''