cuda baby

app.py

@@ -17,7 +17,7 @@ print(
 # Tesla T4
 
 HF_AUTH_TOKEN = "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"
-DEVICE = os.environ.get("cuda:0", "cpu")  # cuda:0
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print("DEVICE SENOOOOOR", DEVICE)
 DTYPE = torch.float32 if DEVICE == "cpu" else torch.float16
 MODEL_NAME = os.environ.get("MODEL_NAME", "NbAiLab/nb-gpt-j-6B-alpaca")