Skip to content

Instantly share code, notes, and snippets.

@rish-16
Last active May 29, 2021 06:15
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save rish-16/9107dfb441c43f770ececfe4107f57ff to your computer and use it in GitHub Desktop.
Save rish-16/9107dfb441c43f770ececfe4107f57ff to your computer and use it in GitHub Desktop.
A guide on Colab TPU training using PyTorch XLA (Part 5)
# add your custom transforms and augmentations
T = transforms.Compose([
transforms.ToTensor(),
...
])
# instructing the master node to download the dataset only ONCE
if not xm.is_master_ordinal():
xm.rendezvous('download_only_once')
mnist_train = datasets.MNIST(..., train=True, transform=T, download=True)
mnist_test = datasets.MNIST(..., train=False, transform=T, download=True)
if xm.is_master_ordinal():
xm.rendezvous('download_only_once')
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment