A guide on Colab TPU training using PyTorch XLA (Part 9)

spawn.py

'''
Configures some pipeline hyper-parameters. You
can set them to whatever you please.

You have the option of either mentioning it here
or creating variables inside the map_fn function.

This is entirely up to you. I do both for demonstration purposes.
'''

flags = {}
flags['batch_size'] = 32
flags['num_workers'] = 8 # we want to train on all 8 cores
flags['num_epochs'] = 10 # I already had the EPOCHS variable in map_fn
flags['seed'] = 42

# start the 8-core TPU and run map_fn on all workers
xmp.spawn(map_fn, args=(flags,), nprocs=8, start_method='fork')