from IPython.core.debugger import set_trace

lr = 0.5  # learning rate
epochs = 2  # how many epochs to train for

for epoch in range(epochs):
    for i in range((n - 1) // bs + 1):
        #         set_trace()
        start_i = i * bs
        end_i = start_i + bs
        [xb](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor") = [x_train](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor")[start_i:end_i]
        [yb](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor") = [y_train](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor")[start_i:end_i]
        [pred](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor") = [model](https://docs.pytorch.org/docs/stable/generated/torch.nn.Sequential.html#torch.nn.Sequential "torch.nn.Sequential")([xb](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor"))
        [loss](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor") = loss_func([pred](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor"), [yb](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor"))

        [loss.backward](https://docs.pytorch.org/docs/stable/generated/torch.Tensor.backward.html#torch.Tensor.backward "torch.Tensor.backward")()
        with [torch.no_grad](https://docs.pytorch.org/docs/stable/generated/torch.no_grad.html#torch.no_grad "torch.no_grad")():
            [weights](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor") -= [weights.grad](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor") * lr
            [bias](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor") -= [bias.grad](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor") * lr
            [weights.grad](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor").zero_()
            [bias.grad](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor").zero_()