refine the device

Author: zxd1997066

test/quantization/quantize_/workflows/int8/test_int8_tensor.py

@@ -30,9 +30,9 @@
 from torchao.testing.model_architectures import ToyTwoLinearModel
 from torchao.testing.utils import TorchAOIntegrationTestCase
 from torchao.utils import (
+    get_available_devices,
     get_current_accelerator_device,
     torch_version_at_least,
-    get_available_devices,
 )
 
 INT8_TEST_CONFIGS = [
@@ -267,15 +267,10 @@ class TestInt8StaticQuant(TorchAOIntegrationTestCase):
     def test_static_activation_per_row_int8_weight(self, granularity, dtype):
         torch.compiler.reset()
 
-<<<<<<< HEAD
         M, N, K = 128, 128, 128
-        input_tensor = torch.randn(M, K, dtype=dtype, device="cuda")
-=======
-        M, N, K = 32, 32, 32
 
         _DEVICE = get_current_accelerator_device()
         input_tensor = torch.randn(M, K, dtype=dtype, device=_DEVICE)
->>>>>>> f07387cd2 (refine the device)
 
         model = torch.nn.Linear(K, N, bias=False).eval().to(device=_DEVICE, dtype=dtype)
         model_static_quant = copy.deepcopy(model)
@@ -353,7 +348,7 @@ def test_static_act_quant_slice_and_select(self, granularity):
         N, K = 256, 512
         M = 32  # batch size
         dtype = torch.bfloat16
-        device = "cuda"
+        device = get_current_accelerator_device()
 
         linear = torch.nn.Linear(K, N, bias=False, dtype=dtype, device=device)
         input_tensor = torch.randn(M, K, dtype=dtype, device=device)
@@ -414,11 +409,12 @@ def test_int8_weight_only_v2_correct_eps(self, dtype):
         torch.manual_seed(42)
 
         # Create test model
-        model = ToyTwoLinearModel(256, 128, 256, dtype=dtype, device="cuda").eval()
+        _DEVICE = get_current_accelerator_device()
+        model = ToyTwoLinearModel(256, 128, 256, dtype=dtype, device=_DEVICE).eval()
         model_baseline = copy.deepcopy(model)
 
         # Create input
-        input_tensor = torch.randn(32, 256, dtype=dtype, device="cuda")
+        input_tensor = torch.randn(32, 256, dtype=dtype, device=_DEVICE)
 
         # Get baseline output
         output_baseline = model_baseline(input_tensor)