add cpu_only

Author: gulsumgudukbay

tests/unit/sharding_compare_test.py

@@ -110,6 +110,7 @@ def compare_sharding_jsons(json1: dict, model1_name: str, json2: dict, model2_na
 
 
 # Requires JAX TPU support to generate the simulated TPU topology.
+@pytest.mark.cpu_only
 @pytest.mark.tpu_backend
 @pytest.mark.parametrize("model_name, topology, num_slice", TEST_CASES)
 def test_sharding_dump_for_model(model_name: str, topology: str, num_slice: str) -> None:
@@ -214,6 +215,7 @@ def abstract_state_and_shardings(request):
   return model_name, topology, num_slice, abstract_state, state_mesh_shardings, logical_shardings
 
 
+@pytest.mark.cpu_only
 @pytest.mark.tpu_backend
 class TestGetAbstractState:
   """Test class for get_abstract_state function and sharding comparison."""