fixed minor tests

Signed-off-by: realAsma <you@example.com>

Author: realAsma

tests/gpu/torch/quantization/test_tensor_quant_cuda.py

@@ -30,14 +30,6 @@
 class TestFakeTensorQuantCuda(FakeTensorQuantTester):
     device = "cuda"
 
-    def test_non_current_gpu(self, need_2_gpus):
-        device = torch.cuda.device_count() - 1
-        assert torch.cuda.current_device() != device
-        x = torch.randn(3, 4).cuda(device)
-        quant_x = tensor_quant.fake_tensor_quant(x, torch.max(torch.abs(x)), None)
-        quant_x_ref = quant(x, torch.max(torch.abs(x)), fake=True)
-        assert torch.allclose(quant_x, quant_x_ref)
-
 
 class TestCudaExt:
     @pytest.mark.parametrize("num_bits", [3, 4, 5, 7, 8, 11])
@@ -145,15 +137,6 @@ def test_backward(self, device):
         loss.backward()
         assert torch.allclose(quant_x.grad, x.grad)
 
-    def test_non_current_gpu(self, need_2_gpus):
-        torch.cuda.set_device(0)
-        device = torch.cuda.device_count() - 1
-        x = torch.randn(3, 4).cuda()
-        quant_x_ref = tensor_quant.fp8_eager(x, torch.tensor(448.0, device=x.device))
-        x = x.cuda(device)
-        quant_x = tensor_quant.scaled_e4m3(x, None, None, 4, 3)
-        assert torch.allclose(quant_x.cuda(), quant_x_ref)
-
     @pytest.mark.parametrize("axis", [0, 1, 2])
     def test_e4m3_per_channel(self, axis):
         x = torch.randn(4, 4, 4, dtype=torch.float32).cuda()

tests/gpu/torch/quantization/test_tensor_quantizer_cuda.py

@@ -55,6 +55,14 @@ def test_e4m3(self, E, M, axis):  # noqa: N803
             ref = tensor_quant.scaled_e4m3(x, e4m3_quantizer._get_amax(x), None, E, M)
             assert torch.allclose(e4m3_x, ref)
 
+    def test_non_current_gpu(self, need_2_gpus):
+        x = torch.randn(3, 4)
+        e4m3_desc = QuantizerAttributeConfig(num_bits=(4, 3), axis=None)
+        quantizer = tensor_quantizer.TensorQuantizer(e4m3_desc).cuda()
+        xq_ref = quantizer(x.to("cuda:0"))
+        xq_test = quantizer(x.to("cuda:1"))
+        assert torch.allclose(xq_ref, xq_test.to("cuda:0"))
+
 
 @pytest.mark.skipif(get_cuda_ext_mx() is None, reason="cuda_ext_mx is not available")
 class TestTensorQuantizerfp4: