peder commited on
Commit
ba319d5
·
1 Parent(s): 8bc3df2

printing gpu

Browse files
Files changed (1) hide show
  1. app.py +6 -0
app.py CHANGED
@@ -8,6 +8,12 @@ import torch
8
  from transformers import pipeline, set_seed
9
  from transformers import AutoTokenizer, AutoModelForCausalLM
10
 
 
 
 
 
 
 
11
 
12
  HF_AUTH_TOKEN = "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"
13
  DEVICE = os.environ.get("cuda:0", "cpu") # cuda:0
 
8
  from transformers import pipeline, set_seed
9
  from transformers import AutoTokenizer, AutoModelForCausalLM
10
 
11
+ import torch
12
+ print(f"Is CUDA available: {torch.cuda.is_available()}")
13
+ # True
14
+ print(
15
+ f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
16
+ # Tesla T4
17
 
18
  HF_AUTH_TOKEN = "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"
19
  DEVICE = os.environ.get("cuda:0", "cpu") # cuda:0