lzhengning
/
JittorMirror
mirror of https://github.com/Jittor/Jittor
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix setitem

This commit is contained in:
li-xl 2020-10-16 20:43:37 +08:00
parent d31cdab30f
commit 84c10c7a06
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
python/jittor/contrib.py
View File

@ -188,7 +188,7 @@ def setitem(x, slices, value):
if isinstance(slices, jt.Var) and slices.dtype == "bool": if isinstance(slices, jt.Var) and slices.dtype == "bool":
mask = jt.broadcast(slices, x) mask = jt.broadcast(slices, x)
value = jt.broadcast(value, x) value = jt.broadcast(value, x)
return x.assign(mask.ternary(value, mask)) return x.assign(mask.ternary(value, x))
if isinstance(slices, list): if isinstance(slices, list):
slices = tuple(slices) slices = tuple(slices)
return x.assign(x.setitem(slices, value)) return x.assign(x.setitem(slices, value))