Update on "[FSDP] Use correct handle training state when prefetching"

[ghstack-poisoned]
pytorch · Apr 4, 2023 · 58bfb34 · 58bfb34
2 parents 971d2c1 + 3c22b26
commit 58bfb34
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/test/distributed/fsdp/test_fsdp_core.py b/test/distributed/fsdp/test_fsdp_core.py
@@ -482,6 +482,7 @@ def _use_unsharded_views_assert_as_tensors(
             *[nn.Linear(3, 3, device=device) for _ in range(NUM_LINEARS)]
         )
         fsdp_model = FSDP(model, **fsdp_kwargs)
+        self.assertEqual(len(list(FSDP.fsdp_modules(fsdp_model))), NUM_LINEARS + 1)
         for _ in range(3):
             inp = torch.randn((2, 3), device=device)
             with self._patch_use_unsharded_views(