So the problem is probably quite obvious. It has to do with the direction of the DAG of gradients.

Namnlös teckning960×720 19.5 KB

A possible way to rework this is
import torch

func = torch.nn.Linear(5, 6)
x = torch.rand(3, 5, requires_grad=True)
x_rows = x.unbind(0)
x = torch.vstack(x_rows)
y_rows = func(x).unbind(0)
for x_row, y_row in zip(x_rows, y_rows):
    torch.autograd.grad(y_row, [x_row], grad_outputs=torch.ones_like(y_row), create_graph=True)[0]

However, I wouldn’t consider this to be a solution as this depends on having access to what is put into the model. The best outcome would be if the experimental feature vmap would turn out to work together with autograd.grad.