fix torch import

Signed-off-by: Vibhu Jawa <[email protected]>
rapidsai · Jul 31, 2024 · 1b02e4b · 1b02e4b
1 parent 91dd34d
commit 1b02e4b
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/test_torch_utils.py b/tests/test_torch_utils.py
@@ -1,8 +1,9 @@
 import pytest
-import torch
 
 from crossfit.utils.torch_utils import concat_and_pad_tensors, pad_tensors
 
+torch = pytest.importorskip("torch")
+
 
 def test_pad_tensors_2d():
     # Test with 2D tensors