Your code look good. It won’t return a tensor of the shape [batch_size, out_dim], so I’ve changed your Module a bit to give the desired output.
class MyLinear(nn.Module):
    def __init__(self, in_dim, h, out_dim):
        super(MyLinear, self).__init__()
        self.weight1 = nn.Parameter(torch.ones(out_dim, h))
        self.weight2= nn.Parameter(torch.randn(h, in_dim))

    def forward(self, x):
        return F.linear(x, torch.mm(self.weight1,self.weight2))

lin = MyLinear(in_dim=10, h=4, out_dim=2)

x = torch.randn(4, 10)
output = lin(x)
output.mean().backward()
lin.weight1.grad
lin.weight2.grad

The output and gradients look good. Could you give a small code example resulting in the NaN loss?