diff --git a/functorch/experimental/__init__.py b/functorch/experimental/__init__.py index 3941f6d96e1f..0500fc2c29d3 100644 --- a/functorch/experimental/__init__.py +++ b/functorch/experimental/__init__.py @@ -1,5 +1,5 @@ # PyTorch forward-mode is not mature yet -from functorch import functionalize from torch._functorch.apis import chunk_vmap from torch._functorch.batch_norm_replacement import replace_all_batch_norm_modules_ from torch._functorch.eager_transforms import hessian, jacfwd, jvp +from torch.func import functionalize