From 2db615ab30dce399bc54251efba7efb16f4e4fed Mon Sep 17 00:00:00 2001 From: Tanima Dey Date: Tue, 7 Oct 2025 14:21:38 -0700 Subject: [PATCH] Make the test case more generic for a given world-size --- .../distributed/_composable/fsdp/test_fully_shard_state_dict.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/distributed/_composable/fsdp/test_fully_shard_state_dict.py b/test/distributed/_composable/fsdp/test_fully_shard_state_dict.py index 2b99bebe51c5b..f5d0badacba27 100644 --- a/test/distributed/_composable/fsdp/test_fully_shard_state_dict.py +++ b/test/distributed/_composable/fsdp/test_fully_shard_state_dict.py @@ -177,7 +177,7 @@ def test_dp_state_dict_cpu_offload(self): def _test_dp_state_dict_cpu_offload( self, offload_policy: CPUOffloadPolicy, cpu_state_dict: bool ): - mlp_dim = 4 + mlp_dim = self.world_size torch.manual_seed(42) with torch.device("meta"): model = nn.Sequential(