Skip to content

Instantly share code, notes, and snippets.

@wookim3

wookim3/DistributedDataParallel.py

Created October 7, 2020 15:35
Show Gist options
  • Save wookim3/a7f461b63cb503d609268bb9c01879cd to your computer and use it in GitHub Desktop.
Save wookim3/a7f461b63cb503d609268bb9c01879cd to your computer and use it in GitHub Desktop.
Download ZIP
Raw
DistributedDataParallel.py
def __init__(self) -> None:
super().__init__()
self.model = self._load_model(pretrained_backbone=True).cuda()
self.model = DistributedDataParallel(self.model, device_ids=[torch.cuda.current_device()],
output_device=torch.cuda.current_device())
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment