autograd with `is_grads_batched=True` fails on GroupNorm about pytorch HOT 4 OPEN

Does it work if you use vmap instead, by any chance

from pytorch.

I'm sorry but I don't know how to use vmap instead;
I tried to vmap the autograd call, but, as stated by https://pytorch.org/tutorials/intermediate/jacobians_hessians.html << We can’t directly apply vmap to torch.autograd.grad; instead, PyTorch provides a torch.func.vjp>>

But apparently torch.func.vjp is designed to build the Jacobian but it runs the forward too.
Whereas my function is non deterministic, so I don't want to (re)run the forward to compute this autograd.grad.

Also, I tried both Cpu and Cuda,
BatchNorm and InstanceNorm work,
Only GroupNorm fails

from pytorch.

Is it possible to avoid running forward the first time? vjp returns both the forward output and the gradients.

from pytorch.

But apparently torch.func.vjp is designed to build the Jacobian but it runs the forward too.
Whereas my function is non deterministic, so I don't want to (re)run the forward to compute this autograd.grad.

You can use torch.func.vjp to run the forward once. It gives you a callable that can be used to run the backward. The idea is to vmap over said callable (this avoids re-running the forward).

from pytorch.