only run torchvision sanity check command test if GPU is present

jfgrimm · Feb 28, 2024 · 454c5ee · 454c5ee
1 parent a2864e7
commit 454c5ee
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/easybuild/easyblocks/t/torchvision.py b/easybuild/easyblocks/t/torchvision.py
@@ -90,11 +90,12 @@ def sanity_check_step(self):
         #      the-cuda-backend-this-could-be-because-the-operator-doesnt-exist-for-this-backend/132352/4
         if self.with_cuda:
             custom_commands = []
-            python_code = '; '.join([
+            python_code = '\n'.join([
                 "import torch, torchvision",
-                "boxes = torch.tensor([[0., 1., 2., 3.]]).to('cuda')",
-                "scores = torch.randn(1).to('cuda')",
-                "print(torchvision.ops.nms(boxes, scores, 0.5))",
+                "if torch.cuda.device_count():",
+                "    boxes = torch.tensor([[0., 1., 2., 3.]]).to('cuda')",
+                "    scores = torch.randn(1).to('cuda')",
+                "    print(torchvision.ops.nms(boxes, scores, 0.5))",
             ])
             custom_commands.append('python -c "%s"' % python_code)
             custom_paths = {