Function torch_gradient_generator

numpy_ml/tests/nn_torch_models.py:21–29

(fn, **kwargs)

def torch_gradient_generator(fn, **kwargs):
    def get_grad(z):
        z1 = torch.autograd.Variable(torch.FloatTensor(z), requires_grad=True)
        z2 = fn(z1, **kwargs).sum()
        z2.backward()
        grad = z1.grad.numpy()
        return grad

    return get_grad


def torch_xe_grad(y, z):