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

model.py

class MyCustomNet(nn.Module):
  def __init__(self, myparams):
    super().__init__()
    
    # define layers
    ...
    
  def forward(self, x):
    '''
    Pass your inputs through your layers as you normally would.
    No change here.
    '''
    return layers(x)