Cannot import torch.func in M1 Pro Mac

Mac OS X
1

I am trying to import vmap and jacrev from torch.func in my MacBook M1 Pro. Here’s what I am doing:

from torch.func import vmap, jacrev

Here’s the error I am getting:

ModuleNotFoundError: No module named 'torch.func'

I am using a conda environment with the latest version of PyTorch and Python 3.8.

How can I resolve this?

2

The import works for me using 2.1.0.dev20230817+cu121.

3

The import did not work for me in a Python or conda virtualenv in the VS code terminal, but works just fine in my MacOS terminal - strange. Thanks for your help.

4

I guess you are using different environments in VS Code vs. your terminal, where one might contain a PyTorch installation which is too old.

1 Like