Hello,
I am trying to write a Function for torch.trace (seems not to exist so far). Here is my code:
class CustomTrace(torch.autograd.Function):
    
    def forward(self, mat):
        self.save_for_backward(mat)
        return torch.Tensor([torch.trace(mat)])
    
    def backward(self, g):
        mat, = self.saved_tensors
        
        return torch.mul(torch.eye(int(mat.size()[0])),g[0])

When I try to check my result, its precision is not optimal, i.e.
input = (torch.autograd.Variable(torch.randn(2,2).double(), requires_grad = True),)
torch.autograd.gradcheck(CustomTrace(), input, eps = 1e-4, atol = 1e-3)

returns true. However,
input = (torch.autograd.Variable(torch.randn(2,2).double(), requires_grad = True),)
torch.autograd.gradcheck(CustomTrace(), input, eps = 1e-6, atol = 1e-4)

returns false. I would hope that my math is right. Is there any problem with my code or is this for some reason to be expected?
Thanks a lot!