Update fully_sharded_data_parallel to fix typing (#110545)

Fixes typing so that linter does not complain when using CustomPolicy.
Pull Request resolved: #110545
Approved by: https://github.com/awgu, https://github.com/Skylion007

Author: mvpatel2000

torch/distributed/fsdp/fully_sharded_data_parallel.py

@@ -99,7 +99,7 @@
     _unshard_params,
     _unshard_params_recurse,
 )
-from .wrap import ModuleWrapPolicy
+from .wrap import CustomPolicy, ModuleWrapPolicy
 
 
 __all__ = [
@@ -261,7 +261,7 @@ class FullyShardedDataParallel(nn.Module, _FSDPState):
             This configures CPU offloading. If this is set to ``None``, then
             no CPU offloading happens. See :class:`CPUOffload` for details.
             (Default: ``None``)
-        auto_wrap_policy (Optional[Union[Callable[[nn.Module, bool, int], bool], ModuleWrapPolicy]]):
+        auto_wrap_policy (Optional[Union[Callable[[nn.Module, bool, int], bool], ModuleWrapPolicy, CustomPolicy]]):
             This specifies a policy to apply FSDP to submodules of ``module``,
             which is needed for communication and computation overlap and thus
             affects performance. If ``None``, then FSDP only applies to
@@ -411,7 +411,9 @@ def __init__(
         process_group: ProcessGroupType = None,
         sharding_strategy: Optional[ShardingStrategy] = None,
         cpu_offload: Optional[CPUOffload] = None,
-        auto_wrap_policy: Optional[Union[Callable, ModuleWrapPolicy]] = None,
+        auto_wrap_policy: Optional[
+            Union[Callable, ModuleWrapPolicy, CustomPolicy]
+        ] = None,
         backward_prefetch: Optional[BackwardPrefetch] = BackwardPrefetch.BACKWARD_PRE,
         mixed_precision: Optional[MixedPrecision] = None,
         ignored_modules: Optional[Iterable[torch.nn.Module]] = None,