Update app.py

app.py

@@ -2,6 +2,7 @@ import random
 import requests
 from bs4 import BeautifulSoup
 from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
 !huggingface-cli login
 
@@ -98,6 +99,10 @@ model_name = 'mistralai/Mistral-7B-Instruct-v0.3'
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
+# Check if a GPU is available and if not, fall back to CPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
 # Example usage
 search_term = "How did Tesla perform in Q1 2024"
 search_results = google_search(search_term, num_results=3)