tennant commited on
Commit
c4529e7
·
verified ·
1 Parent(s): 4d4d201

Update demo.py

Browse files
Files changed (1) hide show
  1. demo.py +2 -1
demo.py CHANGED
@@ -21,7 +21,8 @@ from PIL import Image
21
  transform = vae_transforms('test')
22
  print(f"Is CUDA available: {torch.cuda.is_available()}")
23
  # True
24
- print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
 
25
  # Tesla T4
26
 
27
 
 
21
  transform = vae_transforms('test')
22
  print(f"Is CUDA available: {torch.cuda.is_available()}")
23
  # True
24
+ if device == 'cuda':
25
+ print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
26
  # Tesla T4
27
 
28