[FSDP][Easy] Allow ModuleWrapPolicy to take Iterable

ghstack-source-id: 6e56f5c32da86336e4a943a49f237d448c532cad Pull Request resolved: pytorch#104999
awgu · Jul 21, 2023 · 0159d77 · 0159d77
1 parent 76ebc08
commit 0159d77
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/torch/distributed/fsdp/wrap.py b/torch/distributed/fsdp/wrap.py
@@ -198,13 +198,14 @@ def _module_wrap_policy(
 class ModuleWrapPolicy(_FSDPPolicy):
     """This is a wrapper around :func:`_module_wrap_policy`."""
 
-    def __init__(self, module_classes: Set[Type[nn.Module]]):
+    def __init__(self, module_classes: Iterable[Type[nn.Module]]):
+        module_classes_set = set(module_classes)
         self._policy: Callable = functools.partial(
             _module_wrap_policy,
-            module_classes=module_classes,
+            module_classes=module_classes_set,
         )
-        self._module_classes = module_classes
-        self._module_classes_str = str(module_classes)
+        self._module_classes = module_classes_set
+        self._module_classes_str = str(module_classes_set)
 
     @property
     def policy(self):