I have a following situation. An input has shape [BATCH_SIZE, DIMENSIONALITY] and an output has shape [BATCH_SIZE, CLASSES]. The gradient of the output with respect to the input should have shape [BATCH_SIZE, CLASSES, DIMENSIONALITY] (or similar). Currently, I use a following function to get the gradient:
def compute_grad(input, model):
    input.requires_grad_(True)
    grad = []
    output = model(input)
    for i in range(output.size(1)):
        grad.append(
            torch.autograd.grad(
                outputs=output[:, i],
                inputs=input,
                grad_outputs=torch.ones(len(input)),
                retain_graph=True, 
                create_graph=True
            )[0]
        )
    input.requires_grad_(False)
    return torch.stack(grad)

Is there a way to get the gradient without using for loop? Could this function be written better (more elegant and efficient)?