torch cuda is available

rijobro · rijobro · commit e02a67ce4f94 · 2021-03-22T21:42:33.000Z
Signed-off-by: Richard Brown &lt;33289025+rijobro@users.noreply.github.com&gt;
diff --git a/tests/test_testtimeaugmentation.py b/tests/test_testtimeaugmentation.py
@@ -15,7 +15,6 @@
 
 import numpy as np
 import torch
-from torch._C import has_cuda
 
 from monai.data import CacheDataset, DataLoader, create_test_image_2d
 from monai.data.test_time_augmentation import TestTimeAugmentation
@@ -55,7 +54,7 @@ def tearDown(self) -> None:
 
     def test_test_time_augmentation(self):
         input_size = (20, 20)
-        device = "cuda" if has_cuda else "cpu"
+        device = "cuda" if torch.cuda.is_available() else "cpu"
         num_training_ims = 10
         data = []
         custom_create_test_image_2d = partial(
