# Training loop
for epoch in range(201):
    # Clear gradients
    optimizer.zero_grad()