Please help me diagnose why I'm getting different values for autograd and grads. Here's my code:

import torch

class MLP:
    def __init__(
        self,
        linear_1_in_features,
        linear_1_out_features,
        f_function,
        linear_2_in_features,
        linear_2_out_features,
        g_function
    ):
        """
        Args:
            linear_1_in_features: the in features of first linear layer
            linear_1_out_features: the out features of first linear layer
            linear_2_in_features: the in features of second linear layer
            linear_2_out_features: the out features of second linear layer
            f_function: string for the f function: relu | sigmoid | identity
            g_function: string for the g function: relu | sigmoid | identity
        """
        self.f_function = f_function
        self.g_function = g_function

        self.parameters = dict(
            W1 = torch.randn(linear_1_out_features, linear_1_in_features),
            b1 = torch.randn(linear_1_out_features),
            W2 = torch.randn(linear_2_out_features, linear_2_in_features),
            b2 = torch.randn(linear_2_out_features),
        )
        self.grads = dict(
            dJdW1 = torch.zeros(linear_1_out_features, linear_1_in_features),
            dJdb1 = torch.zeros(linear_1_out_features),
            dJdW2 = torch.zeros(linear_2_out_features, linear_2_in_features),
            dJdb2 = torch.zeros(linear_2_out_features),
        )

        # put all the cache value you need in self.cache
        self.cache = dict()

    def forward(self, x):
        """
        Args:
            x: tensor shape (batch_size, linear_1_in_features)
        """
        # TODO: Implement the forward function
        self.cache['x'] = x

        z1 = self.cache['z1'] = x @ self.parameters['W1'].T + self.parameters['b1']

        if self.f_function == 'relu':
            a1 = self.cache['a1'] = torch.relu(z1)
        elif self.f_function == 'sigmoid':
            a1 = self.cache['a1'] = torch.sigmoid(z1)
        else:
            a1 = self.cache['a1'] = z1

        z2 = self.cache['z2'] = a1 @ self.parameters['W2'].T + self.parameters['b2']

        if self.g_function == 'relu':
            a2 = self.cache['a2'] = torch.relu(z2)
        elif self.g_function == 'sigmoid':
            a2 = self.cache['a2'] = torch.sigmoid(z2)
        else:
            a2 = self.cache['a2'] = z2

        return a2
    
    def backward(self, dJdy_hat):
        """
        Args:
            dJdy_hat: The gradient tensor of shape (batch_size, linear_2_out_features)
        """
        # TODO: Implement the backward function
        x = self.cache['x']
        a1 = self.cache['a1']
        z1 = self.cache['z1']
        z2 = self.cache['z2']

        if self.g_function == 'relu':
            dJdz2 = dJdy_hat * (z2 > 0).float()
        elif self.g_function == 'sigmoid':
            sigmoid_z2 = torch.sigmoid(z2)
            dJdz2 = dJdy_hat * sigmoid_z2 * (1 - sigmoid_z2)
        else:
            dJdz2 = dJdy_hat

        self.grads['dJdW2'] = dJdz2.T @ a1
        self.grads['dJdb2'] = dJdz2.sum(dim=0)

        dJda1 = dJdz2 @ self.parameters["W2"]
        
        if self.f_function == 'relu':
            dJdz1 = dJda1 * (z1 > 0).float()
        elif self.f_function == 'sigmoid':
            sigmoid_z1 = torch.sigmoid(z1)
            dJdz1 = dJda1 * sigmoid_z1 * (1 - sigmoid_z1)
        else:
            dJdz1 = dJda1
        
        self.grads['dJdW1'] = dJdz1.T @ x
        self.grads['dJdb1'] = dJdz1.sum(dim=0)

        
    def clear_grad_and_cache(self):
        for grad in self.grads:
            self.grads[grad].zero_()
        self.cache = dict()

def mse_loss(y, y_hat):
    """
    Args:
        y: the label tensor (batch_size, linear_2_out_features)
        y_hat: the prediction tensor (batch_size, linear_2_out_features)

    Return:
        J: scalar of loss
        dJdy_hat: The gradient tensor of shape (batch_size, linear_2_out_features)
    """
    # TODO: Implement the mse loss
    N = y.shape[0]
    J = torch.sum((y - y_hat) ** 2) / N
    dJdy_hat = -2 * (y - y_hat) / N
    return J, dJdy_hat
    

def bce_loss(y, y_hat):
    """
    Args:
        y_hat: the prediction tensor
        y: the label tensor
        
    Return:
        loss: scalar of loss
        dJdy_hat: The gradient tensor of shape (batch_size, linear_2_out_features)
    """
    # TODO: Implement the bce loss
    N = y.shape[0]
    loss = -torch.sum(y * torch.log(y_hat) + (1 - y) * torch.log(1 - y_hat)) / N
    dJdy_hat = -(y / y_hat - (1 - y) / (1 - y_hat)) / N
    return loss, dJdy_hat

Here's the test file:

from collections import OrderedDict

import torch
import torch.nn as nn
import torch.nn.functional as F
from mlp import MLP, mse_loss, bce_loss

net = MLP(
    linear_1_in_features=2,
    linear_1_out_features=20,
    f_function='relu',
    linear_2_in_features=20,
    linear_2_out_features=5,
    g_function='identity'
)
x = torch.randn(10, 2)
y = torch.randn(10, 5)

net.clear_grad_and_cache()
y_hat = net.forward(x)
J, dJdy_hat = mse_loss(y, y_hat)
net.backward(dJdy_hat)

#------------------------------------------------
# compare the result with autograd
net_autograd = nn.Sequential(
    OrderedDict([
        ('linear1', nn.Linear(2, 20)),
        ('relu', nn.ReLU()),
        ('linear2', nn.Linear(20, 5)),
    ])
)
net_autograd.linear1.weight.data = net.parameters['W1']
net_autograd.linear1.bias.data = net.parameters['b1']
net_autograd.linear2.weight.data = net.parameters['W2']
net_autograd.linear2.bias.data = net.parameters['b2']

y_hat_autograd = net_autograd(x)
y_hat_autograd.retain_grad()

J_autograd = F.mse_loss(y_hat_autograd, y)

net_autograd.zero_grad()
J_autograd.backward()

print(net_autograd.linear1.weight.grad.data)
print(net.grads['dJdW1'])

print((net_autograd.linear1.weight.grad.data - net.grads['dJdW1']).norm() < 1e-3)
print((net_autograd.linear1.bias.grad.data - net.grads['dJdb1']).norm() < 1e-3)
print((net_autograd.linear2.weight.grad.data - net.grads['dJdW2']).norm() < 1e-3)
print((net_autograd.linear2.bias.grad.data - net.grads['dJdb2']).norm()< 1e-3)
#------------------------------------------------

And Here's the respective tensors for autograd and grads:

tensor([[ 2.4838, -4.9587],
        [-5.0845, -0.4971],
        [ 0.5651, -1.4917],
        [-0.3781,  1.3107],
        [ 0.0336,  2.0604],
        [-0.7053, -3.0871],
        [ 0.7716, -0.6366],
        [-1.4441,  1.3391],
        [ 0.6104, -0.6521],
        [-0.2280,  2.1847],
        [ 0.1574, -2.3221],
        [-4.7131,  1.2860],
        [-0.8992, -1.0204],
        [-0.3116,  0.7592],
        [ 0.6709, -1.3290],
        [-2.3996,  4.0735],
        [ 1.7576, -4.8977],
        [ 0.5142, -0.8679],
        [-5.6518,  4.4894],
        [-0.0110, -0.0226]])
tensor([[ 12.4190, -24.7933],
        [-25.4227,  -2.4856],
        [  2.8255,  -7.4584],
        [ -1.8905,   6.5534],
        [  0.1679,  10.3022],
        [ -3.5263, -15.4357],
        [  3.8582,  -3.1828],
        [ -7.2204,   6.6955],
        [  3.0521,  -3.2605],
        [ -1.1402,  10.9235],
        [  0.7868, -11.6103],
        [-23.5653,   6.4299],
        [ -4.4958,  -5.1022],
        [ -1.5580,   3.7958],
        [  3.3546,  -6.6452],
        [-11.9982,  20.3676],
        [  8.7882, -24.4887],
        [  2.5711,  -4.3393],
        [-28.2591,  22.4469],
        [ -0.0549,  -0.1128]])