File size: 280 Bytes
375a1cf
 
 
 
 
 
1
2
3
4
5
6
7
# PyTorch forward-mode is not mature yet
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 functorch import functionalize