AttributeError: module 'torch' has no attribute 'vmap'

I am running pytorch version 1.13 on a Google Cloud instance. When I try to use torch.vmap, I get the attribute error. I know this was an issue in older versions of pytorch, but not sure why it isn’t working with my version.

Best,
David

torch.vmap was added in 2.0 so you might need to use functorch.vmap in older releases or update PyTorch to the latest stable or nightly release.