Fix ManagedCollisionEmbeddingCollection not recognised as unpooled if…

… it is a submodule

Summary: When `ManagedCollisionEmbeddingCollection` is part of a module's submodules, the parent module is not recognised as unpooled. Doesn't matter in practice, since EC is always a submodule of MC_EC, but it is more correct this way

Reviewed By: henrylhtsang

Differential Revision: D53920749

torchrec/distributed/planner/types.py

@@ -371,7 +371,7 @@ def module_pooled(module: nn.Module, sharding_option_name: str) -> bool:
 
         for submodule in module.modules():
             if isinstance(submodule, EmbeddingCollectionInterface) or isinstance(
-                module, ManagedCollisionEmbeddingCollection
+                submodule, ManagedCollisionEmbeddingCollection
             ):
                 for name, _ in submodule.named_parameters():
                     if sharding_option_name in name: