6 lines
273 B
Python
6 lines
273 B
Python
|
# 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
|