calculating-gradient.py

def calc_grad(xb, yb, model):
    preds = model(xb)
    loss = mnist_loss(preds, yb)
    loss.backward()
weights.grad.zero_()
bias.grad.zero_();