taichi
taichi copied to clipboard
ti.ad.no_grad decorator not support return
Describe the bug If I using ti.ad.no_grad decorator the returns replaces to None
To Reproduce
import torch
device = 'cuda'
import taichi as ti
ti.init(getattr(ti,device))
ti2d = ti.types.ndarray(ndim=2, needs_grad=False)
@ti.kernel
def taichi_sinsum (x: ti2d) -> ti.f32:
s = 0.
for i,j in x:
s += ti.sin(x[i,j])
return s
@ti.ad.no_grad
def taichi_nograd(x):
s = taichi_sinsum (x)
print(f"in nograd func s={s}")
return s
N = 64
x = torch.rand([N,N],device=device,requires_grad=True)
s = taichi_nograd(x)
print(f"in main s={s}")
Log/Screenshots
[Taichi] version 1.8.0, llvm 15.0.1, commit 5a994612, win, python 3.8.8
[I 01/10/24 12:06:14.158 13504] [shell.py:_shell_pop_print@23] Graphical python shell detected, using wrapped sys.stdout
[Taichi] Starting on arch=cuda
in nograd func s=1868.820068359375
in main s=None
Additional comments
If possible, please also consider attaching the output of command ti diagnose. This produces the detailed environment information and hopefully helps us diagnose faster.
If you have local commits (e.g. compile fixes before you reproduce the bug), please make sure you first make a PR to fix the build errors and then report the bug.