frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[functorch] Update eager_fusion.py

Browse Source
This commit is contained in:
Horace He
2021-08-01 16:53:07 -07:00
committed by Jon Janzen
parent 02713fd157
commit 5466ac5da3
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
functorch/examples/compilation/eager_fusion.py
View File

@ -11,10 +11,10 @@ def f(a):
fw_compiler = partial(tvm_compile, name='fw_keops')
bw_compiler = partial(tvm_compile, name='bw_keops')
# fw_compiler = lambda x, _: x
# bw_compiler = lambda x, _: x
compiled_f = compiled_function(f, fw_compiler, bw_compiler).apply
# fw_compiler = lambda x, _: x
# bw_compiler = lambda x, _: x
iters = 10
out = compiled_f(a)
out.sum().backward()