Skip to content

Instantly share code, notes, and snippets.

@ViktorAlm
ViktorAlm / ddp_example.py
Created October 8, 2021 15:08 — forked from sgraaf/ddp_example.py
PyTorch Distributed Data Parallel (DDP) example
#!/usr/bin/env python
# -*- coding: utf-8 -*-
from argparse import ArgumentParser
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, Dataset
from torch.utils.data.distributed import DistributedSampler
from transformers import BertForMaskedLM