File size: 273 Bytes
d1ceb73 |
1 2 3 4 5 6 |
# 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
|