Update app.py

app.py

@@ -15,13 +15,23 @@ from transformers import AutoTokenizer, AutoModel
 hf_token = os.getenv("HF_TOKEN_READ_WRITE")  # Read the token from Secrets
 login(hf_token)
 
+if torch.cuda.is_available():
+    print("✅ GPU is available")
+    print("GPU Name:", torch.cuda.get_device_name(0))
+else:
+    print("❌ No GPU available")
+
 # ---------------------------------------------------------------------------
 # 1. Define model name and load model/tokenizer
 # ---------------------------------------------------------------------------
 model_name = "mistralai/Mistral-7B-Instruct-v0.3"  # fictional placeholder
 
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)  
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = AutoModelForCausalLM.from_pretrained(model_name, token=hf_token, torch_dtype=torch.float16, device_map="auto").to(device)
+
+print(f"✅ Model loaded on {device}")
+#model = AutoModelForCausalLM.from_pretrained(model_name)  
 
 # ---------------------------------------------------------------------------
 # 2. Define a tiny "dataset" for demonstration