update ut

PaddlePaddle · Jan 10, 2024 · 91777b2 · 91777b2
1 parent fe580a3
commit 91777b2
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/test/auto_parallel/test_shard_tensor_api.py b/test/auto_parallel/test_shard_tensor_api.py
@@ -93,9 +93,9 @@ def test_dynamic_mode_property_change(self):
         self.assertEqual(d_tensor.process_mesh, self.mesh)
 
     def test_stop_gradient(self):
-        x = paddle.ones([10, 10])
+        x = paddle.ones([4, 1024, 512])
         x.stop_gradient = False
-        x = dist.shard_tensor(x, self.mesh, [Shard(0)])
+        x = dist.shard_tensor(x, self.mesh, [Shard(0), Replicate()])
         assert not x.stop_gradient