Skip to content

Instantly share code, notes, and snippets.

View revantteotia's full-sized avatar

Revant Teotia revantteotia

View GitHub Profile
@revantteotia
revantteotia / ddp_example.py
Created July 12, 2021 06:47 — 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