Unverified Commit 7a062894 authored by Stas Bekman's avatar Stas Bekman Committed by GitHub
Browse files

[doc] make the basic example more "compilable" (#207)

* make the basic example usable out of the box
* clarify
parent 2b121242
...@@ -54,6 +54,7 @@ See a more complete example [here](https://github.com/facebookresearch/fairscale ...@@ -54,6 +54,7 @@ See a more complete example [here](https://github.com/facebookresearch/fairscale
```python ```python
import torch import torch
import torch.distributed as dist
import torch.multiprocessing as mp import torch.multiprocessing as mp
from fairscale.optim.oss import OSS from fairscale.optim.oss import OSS
from torch.nn.parallel import DistributedDataParallel as DDP from torch.nn.parallel import DistributedDataParallel as DDP
...@@ -63,8 +64,8 @@ def train( ...@@ -63,8 +64,8 @@ def train(
world_size: int, world_size: int,
epochs: int): epochs: int):
# DDP # DDP init example
dist_init(rank, world_size) dist.init_process_group(backend='nccl', init_method="tcp://localhost:29501", rank=rank, world_size=world_size)
# Problem statement # Problem statement
model = myAwesomeModel().to(rank) model = myAwesomeModel().to(rank)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment