Move test_diloco_two_slices to cpu_only

Charles Li · Charles Li · commit 792ac1946adc · 2026-04-28T17:18:54.000Z
diff --git a/tests/unit/diloco_test.py b/tests/unit/diloco_test.py
@@ -287,7 +287,8 @@ def test_diloco_qwen3_moe_two_slices(self):
         )
     )
 
-  @pytest.mark.tpu_only
+  @pytest.mark.cpu_only
+  @pytest.mark.tpu_backend
   def test_diloco_two_slices(self):
     temp_dir = gettempdir()
     compiled_trainstep_file = os.path.join(temp_dir, "test_compiled_diloco.pickle")

Original file line number	Diff line number	Diff line change
`@@ -287,7 +287,8 @@ def test_diloco_qwen3_moe_two_slices(self):`
`287`	`287`	`)`
`288`	`288`	`)`
`289`	`289`
`290`		`- @pytest.mark.tpu_only`
	`290`	`+ @pytest.mark.cpu_only`
	`291`	`+ @pytest.mark.tpu_backend`
`291`	`292`	`def test_diloco_two_slices(self):`
`292`	`293`	`temp_dir = gettempdir()`
`293`	`294`	`compiled_trainstep_file = os.path.join(temp_dir, "test_compiled_diloco.pickle")`