Spaces:

tennant
/

old_tok

Sleeping

tennant commited on Mar 10

Commit

c4529e7

verified ·

1 Parent(s): 4d4d201

Update demo.py

Files changed (1) hide show

demo.py CHANGED Viewed

@@ -21,7 +21,8 @@ from PIL import Image
 transform = vae_transforms('test')
 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

 transform = vae_transforms('test')
 print(f"Is CUDA available: {torch.cuda.is_available()}")
 # True
+if device == 'cuda':
+    print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
 # Tesla T4