|
9 | 9 | import torch |
10 | 10 | import torch.nn.functional as F |
11 | 11 |
|
12 | | -from torchao.testing.utils import skip_if_no_cuda |
13 | 12 | from torchao.utils import TorchAOBaseTensor, torch_version_at_least |
14 | 13 |
|
15 | 14 |
|
@@ -128,14 +127,15 @@ def _test_default_impls_helper(self, lp_tensor, lp_tensor_for_copy): |
128 | 127 | self.assertTrue(torch.equal(lp_tensor.qdata, reconstructed.qdata)) |
129 | 128 | self.assertEqual(lp_tensor.attr, reconstructed.attr) |
130 | 129 |
|
| 130 | + device = torch.accelerator.current_accelerator() |
131 | 131 | # test _get_to_kwargs |
132 | | - _ = lp_tensor._get_to_kwargs(torch.strided, device="cuda") |
133 | | - _ = lp_tensor._get_to_kwargs(layout=torch.strided, device="cuda") |
| 132 | + _ = lp_tensor._get_to_kwargs(torch.strided, device=device) |
| 133 | + _ = lp_tensor._get_to_kwargs(layout=torch.strided, device=device) |
134 | 134 |
|
135 | 135 | # `to` / `_to_copy` |
136 | 136 | original_device = lp_tensor.device |
137 | | - lp_tensor = lp_tensor.to("cuda") |
138 | | - self.assertEqual(lp_tensor.device.type, "cuda") |
| 137 | + lp_tensor = lp_tensor.to(device) |
| 138 | + self.assertEqual(lp_tensor.device.type, device.type) |
139 | 139 | lp_tensor = lp_tensor.to(original_device) |
140 | 140 | self.assertEqual(lp_tensor.device, original_device) |
141 | 141 |
|
@@ -218,7 +218,7 @@ def _test_default_impls_helper(self, lp_tensor, lp_tensor_for_copy): |
218 | 218 | getattr(lp_tensor, tensor_attribute_name), |
219 | 219 | ) |
220 | 220 |
|
221 | | - @skip_if_no_cuda() |
| 221 | + @unittest.skipIf(not torch.accelerator.is_available(), "Need CUDA available") |
222 | 222 | def test_default_impls(self): |
223 | 223 | """Making sure some common functions has default implementations, such as |
224 | 224 | __tensor_unflatten__, __tensor_flatten__, _apply_fn_to_data, __repr__, to |
@@ -248,7 +248,7 @@ def __init__(self, qdata, attr, device): |
248 | 248 | lp_tensor_for_copy = MyTensor(another_tensor, "attr", None) |
249 | 249 | self._test_default_impls_helper(lp_tensor, lp_tensor_for_copy) |
250 | 250 |
|
251 | | - @skip_if_no_cuda() |
| 251 | + @unittest.skipIf(not torch.accelerator.is_available(), "Need GPU available") |
252 | 252 | def test_default_impls_with_optional_data(self): |
253 | 253 | class MyTensorWithOptionalData(TorchAOBaseTensor): |
254 | 254 | tensor_data_names = ["qdata"] |
@@ -285,7 +285,7 @@ def __init__(self, qdata, attr, device, zero_point=None): |
285 | 285 | ) |
286 | 286 | self._test_default_impls_helper(lp_tensor, lp_tensor_for_copy) |
287 | 287 |
|
288 | | - @skip_if_no_cuda() |
| 288 | + @unittest.skipIf(not torch.accelerator.is_available(), "Need GPU available") |
289 | 289 | def test_default_impls_with_optional_attr(self): |
290 | 290 | class MyTensorWithOptionalData(TorchAOBaseTensor): |
291 | 291 | tensor_data_names = ["qdata"] |
|
0 commit comments