Thanks for the links.
For my case of numerical gradients I got to this:
import torch
import numpy as np
from torch.autograd import gradcheck

eps = 1e-6

g_A = torch.rand(3, 3)
t_B = torch.rand(3, 3)
o_C = {'foo': 0, 'bar': [0, 1]}

g_A.requires_grad = True

# something simple that is actually differentiable, but it simulates a non-differentiable function.
# it also has a non tensor input
def loss_nondiff(A, B, C):
    a = A.data.numpy()
    b = B.numpy()[0]
    
    cost = np.expand_dims(np.sum(a)+np.sum(b), 0)
    print(C, cost)
    
    return torch.from_numpy(cost.astype(np.float32))

# autograd wrapper
class loss_test(torch.autograd.Function):
    
    @staticmethod
    def forward(ctx, A, B):
        ctx.save_for_backward(A, B)
        
        return loss_nondiff(A, B, o_C)
    
    @staticmethod
    def backward(ctx, grad_output):
        A, B = ctx.saved_tensors
        
        zeros = torch.from_numpy(np.zeros((3, 3)).astype(np.float32))
        grad_A = torch.from_numpy(np.zeros((3, 3)).astype(np.float32))
        
        for i in range(3):
            for j in range(3):
                teps = zeros.clone()
                teps[i, j] += eps
                
                # this grad_output is just to shut up the gradcheck, since it's always 1 for the loss function
                # i have no ideia what to do with it, specially if has a different size
                grad_A[i, j] = grad_output*(loss_nondiff(A+teps, B, o_C)-loss_nondiff(A-teps, B, o_C))/(2*eps)
        
        return grad_A, None
    

criterion = loss_test.apply
gradcheck(criterion, (g_A, t_B))

This works because it is a simple case, but for the case where the function is not a loss function, and you can have outputs of different sizes and/or multiple outputs, I still don’t know what to do grad_output or with the gradient computations.