1# PyTorch forward-mode is not mature yet 2from functorch import functionalize 3from torch._functorch.apis import chunk_vmap 4from torch._functorch.batch_norm_replacement import replace_all_batch_norm_modules_ 5from torch._functorch.eager_transforms import hessian, jacfwd, jvp 6