train.py

## training
for i in range(150):
    model = model.train()

    ## forward
    output = model(x)
    loss = criterion(output, y)
    optimizer.zero_grad()

    ## backward + update model params 
    loss.backward()
    optimizer.step()

    model.eval()
    print('Epoch: %d | Loss: %.4f' %(i, loss.detach().item()))