Update dynamo test to be less constrain (#6981)

pytorch · Apr 26, 2024 · b9a9449 · b9a9449
1 parent b3be775
commit b9a9449
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/test/dynamo/test_dynamo.py b/test/dynamo/test_dynamo.py
@@ -641,10 +641,7 @@ def test_all_cpu_tensor(self):
       # there should be 18 paramters + 1 input
       self.assertGreater(len(w), 15)
       self.assertIn('Found tensor with shape torch.Size', str(w[0].message))
-    # no XLA operation should happens except a empty mark_step. Partitioner should offload all CPU
-    # ops to CPU.
-    self.assertEqual(len(met.counter_names()), 1)
-    self.assertIn('MarkStep', met.counter_names())
+    self.assertLessEqual(len(met.counter_names()), 1)
 
 
 class DynamoOperationsTests(test_utils.XlaTestCase):