Update app.py

app.py

@@ -3,8 +3,17 @@ from peft import PeftModel, PeftConfig
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
+
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    print("GPU is available!")
+else:
+    device = torch.device("cpu")
+    print("GPU is not available, using CPU.")
+
 # Load the model and config when the script starts
-peft_model_id = "phearion/bigbrain-v0.0.1"
+peft_model_id =
+ "phearion/bigbrain-v0.0.1"
 config = PeftConfig.from_pretrained(peft_model_id)
 model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, return_dict=True)
 tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)