Run MultiNestedTensor tests on both CPU and GPU (#476)

pyg-team · Jan 1, 2025 · 17b5507 · 17b5507
1 parent febb5e4
commit 17b5507
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/test/data/test_multi_nested_tensor.py b/test/data/test_multi_nested_tensor.py
@@ -6,7 +6,7 @@
 from torch import Tensor
 
 from torch_frame.data import MultiNestedTensor
-from torch_frame.testing import onlyCUDA
+from torch_frame.testing import onlyCUDA, withCUDA
 
 
 def assert_equal(tensor_mat: list[list[Tensor]],
@@ -95,7 +95,7 @@ def test_fillna_col():
                           torch.tensor([100], dtype=torch.float32)))
 
 
-@onlyCUDA
+@withCUDA
 def test_basics(device):
     num_rows = 8
     num_cols = 10