disable tests using mark_sharding + assume_pure on GPU (#9013)

Author: qihqi

test/scan/test_scan_spmd.py

@@ -1,3 +1,4 @@
+import os
 import sys
 import re
 import unittest
@@ -230,11 +231,19 @@ def check_dots_in_model(self, model, x, expect_pattern):
   def count_regex(self, hlo_text, regex_str):
     return len(re.findall(regex_str, hlo_text))
 
+  @unittest.skipIf(
+      torch.cuda.is_available() or os.environ.get('PJRT_DEVICE') == 'CUDA',
+      "TODO(https://github.com/pytorch/xla/issues/9017): Get these tests working on GPU"
+  )
   def test_assume_pure_works_with_mark_sharding(self):
     x = torch.randn((3, 4, 5, 128), device='xla')
     assume_pure(mark_sharding)(x, self.spmd_mesh, ("model", None, None, None))
     # assert not throwing
 
+  @unittest.skipIf(
+      torch.cuda.is_available() or os.environ.get('PJRT_DEVICE') == 'CUDA',
+      "TODO(https://github.com/pytorch/xla/issues/9017): Get these tests working on GPU"
+  )
   def test_convert_to_jax_mesh(self):
     jax_mesh = self.spmd_mesh.maybe_convert_and_get_jax_mesh()
     self.assertEqual(jax_mesh.devices.shape, self.spmd_mesh.mesh_shape)