Update main.py

main.py

@@ -1,14 +1,20 @@
 #Imporing required libraries
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
+import torch
 
 # Defining the pipeline and the model
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 tokenizer = AutoTokenizer.from_pretrained("huggyllama/llama-7b")
-model = AutoModelForCausalLM.from_pretrained("huggyllama/llama-7b")
+model = AutoModelForCausalLM.from_pretrained("huggyllama/llama-7b", device_map="auto", torch_dtype=torch.float16)
+
+#model = model.to(device)
 
 print("***")
 print("Loaded tokenizer and model")
+print(device)
 print("***")
 
 pipe_flan = pipeline("text-generation", model=model, tokenizer=tokenizer)