_posts/2022-3-14-introducing-pytorch-fully-sharded-data-parallel-api.md

About this example _posts/2022-3-14-introducing-pytorch-fully-sharded-data-parallel-api.md

from torch.distributed.fsdp import (
   FullyShardedDataParallel,
   CPUOffload,
)
from torch.distributed.fsdp.wrap import (
   enable_wrap,
   wrap,
)
import torch.nn as nn
from typing import Dict
 
 
class model(nn.Module):
   def __init__(self):
       super().__init__()
       self.layer1 = wrap(nn.Linear(8, 4))
       self.layer2 = nn.Linear(4, 16)
       self.layer3 = wrap(nn.Linear(16, 4))
 
wrapper_kwargs = Dict(cpu_offload=CPUOffload(offload_params=True))
with enable_wrap(wrapper_cls=FullyShardedDataParallel, **wrapper_kwargs):
   fsdp_model = wrap(model())

Traceback:

TypeError                                 Traceback (most recent call last)
/var/tmp/ipykernel_54070/554172377.py in <module>
     18        self.layer3 = wrap(nn.Linear(16, 4))
     19 
---> 20 wrapper_kwargs = Dict(cpu_offload=CPUOffload(offload_params=True))
     21 with enable_wrap(wrapper_cls=FullyShardedDataParallel, **wrapper_kwargs):
     22    fsdp_model = wrap(model())

~/miniconda3/envs/py39/lib/python3.9/typing.py in __call__(self, *args, **kwargs)
    678     def __call__(self, *args, **kwargs):
    679         if not self._inst:
--> 680             raise TypeError(f"Type {self._name} cannot be instantiated; "
    681                             f"use {self.__origin__.__name__}() instead")
    682         result = self.__origin__(*args, **kwargs)

TypeError: Type Dict cannot be instantiated; use dict() instead

Open _posts/2022-3-14-introducing-pytorch-fully-sharded-data-parallel-api.md and inspect the Python example around wrapper_kwargs. Run the example or compare it with Python's typing behavior to confirm the shown TypeError is resolved; done means the documented FSDP example runs without that error.