davanstrien HF staff commited on
Commit
faa5dab
·
1 Parent(s): c14bbc6

add cuda check

Browse files
Files changed (1) hide show
  1. app.py +4 -0
app.py CHANGED
@@ -12,6 +12,10 @@ import transformers
12
  from huggingface_hub import CommitScheduler, hf_hub_download, login
13
  from transformers import AutoTokenizer
14
 
 
 
 
 
15
  HF_TOKEN = os.getenv("HF_TOKEN")
16
  login(HF_TOKEN)
17
 
 
12
  from huggingface_hub import CommitScheduler, hf_hub_download, login
13
  from transformers import AutoTokenizer
14
 
15
+ print(f"Is CUDA available: {torch.cuda.is_available()}")
16
+
17
+ print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
18
+
19
  HF_TOKEN = os.getenv("HF_TOKEN")
20
  login(HF_TOKEN)
21