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

model.py

device = xm.xla_device()

# define some hyper-params you'd feed into your model
in_channels = ...
random_param = ...

# create model using appropriate hyper-params
net = MyCustomNet(...)

# seat it atop the TPU worker device and switch it to train mode
net = net.to(device).train()

# get the loss function and optimizer – use anything
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=..., betas=(...))