RuntimeError: vmap: Calling Tensor.as_strided is not supported

autograd
1

I am using torch.autograd.functional.hessian to compute the hessian in my optimization problem, but I keep getting the following error that I’m not finding many resources on:

RuntimeError: vmap: Calling Tensor.as_strided is not supported unless the batch dims being vmapped over are at the front of the tensor (in memory layout). When they are not at the front of the tensor this operation can be error prone so we actively discourage it; please file us a bug report and/or try to express the as_strided operation in terms of PyTorch view operations

My code is somewhat involved so it’s hard for me to post a minimal working example. I don’t even know what to be looking for though that could be causing this issue. Any pointers to help me out? I thought maybe it my usage of view in spots so I replaced those with reshape but it had no effect. I’m not sure where to go from here.

2

as_strided does very wild things with the Tensor, so it’s not surprising that vmap doesn’t like it.
If you are not using as_strided directly, you could see if you can get a traceback that says which function (or backward(?)) causes the error.

Best regards

Thomas

3

Thanks for the suggestion. It turns out we were setting vectorize=True in the hessian function which sounds like it’s doing strange things with vmap and it’s indicated it’s an experimental feature. Once I changed it back to False all was good.