printing gpu

app.py

@@ -8,6 +8,12 @@ import torch
 from transformers import pipeline, set_seed
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
+import torch
+print(f"Is CUDA available: {torch.cuda.is_available()}")
+# True
+print(
+    f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
+# Tesla T4
 
 HF_AUTH_TOKEN = "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"
 DEVICE = os.environ.get("cuda:0", "cpu")  # cuda:0