Update app.py

app.py

@@ -3,6 +3,11 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, StopStringCriteria, StoppingCriteriaList
 import torch
 
+import torch
+import os
+os.system("nvidia-smi")
+print("TORCH_CUDA", torch.cuda.is_available())
+
 # Load the tokenizer and model
 repo_name = "nvidia/Hymba-1.5B-Instruct"
 
@@ -10,6 +15,9 @@ tokenizer = AutoTokenizer.from_pretrained(repo_name, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(repo_name, trust_remote_code=True)
 model = model.cuda().to(torch.bfloat16)
 
+print("model is loaded")
+
+
 # Chat with Hymba
 # prompt = input()
 prompt = "Who are you?"
@@ -22,6 +30,10 @@ messages.append({"role": "user", "content": prompt})
 # Apply chat template
 tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to('cuda')
 stopping_criteria = StoppingCriteriaList([StopStringCriteria(tokenizer=tokenizer, stop_strings="</s>")])
+
+print("generating prompt")
+
+
 outputs = model.generate(
     tokenized_chat, 
     max_new_tokens=256,