check

Lightning-AI · Jan 8, 2025 · 9fb71dc · 9fb71dc
1 parent e7d800f
commit 9fb71dc
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 2 deletions.
diff --git a/.azure/gpu-tests-fabric.yml b/.azure/gpu-tests-fabric.yml
@@ -120,7 +120,7 @@ jobs:
           set -e
           python requirements/collect_env_details.py
           python -c "import torch ; mgpu = torch.cuda.device_count() ; assert mgpu == 2, f'GPU: {mgpu}'"
-          python -c "import bitsandbytes"
+          python requirements/pytorch/check-avail-extras.py
         displayName: "Env details"
 
       - bash: python -m pytest lightning_fabric

diff --git a/.azure/gpu-tests-pytorch.yml b/.azure/gpu-tests-pytorch.yml
@@ -135,7 +135,6 @@ jobs:
           python requirements/collect_env_details.py
           python -c "import torch ; mgpu = torch.cuda.device_count() ; assert mgpu == 2, f'GPU: {mgpu}'"
           python requirements/pytorch/check-avail-extras.py
-          python -c "import bitsandbytes"
         displayName: "Env details"
 
       - bash: python -m pytest pytorch_lightning

diff --git a/requirements/pytorch/check-avail-extras.py b/requirements/pytorch/check-avail-extras.py
@@ -4,3 +4,7 @@
     import matplotlib  # noqa: F401
     import omegaconf  # noqa: F401
     import rich  # noqa: F401
+
+    import torch  # noqa: F401
+    if torch.cuda.is_available():
+        import bitsandbytes  # noqa: F401