Skip to content

Commit 71a3b20

Browse files
committed
Return the updated inp from seteitem_
The following trace ``` def higher_order_autograd_function_apply_brtuu(any5, t_0, t_1): # <eval_with_key>.83:5: _set_grad_enabled = torch._C._set_grad_enabled(False); _set_grad_enabled = None ltorch._set_grad_enabled_with_warning(False) # <eval_with_key>.83:6: output = torch.zeros(4096, 32, 128, device = device(type='cuda', index=0), dtype = torch.bfloat16) t2 = ltorch.zeros(4096, 32, 128, device=devices.Device("cuda:0"), dtype=torch.bfloat16) # t2: "cuda:0 bf16[4096, 32, 128] self._requires_grad=False" # t2 = ltorch.full((4096, 32, 128), 0, device=devices.Device("cuda:0"), dtype=torch.bfloat16) # t2: "cuda:0 bf16[4096, 32, 128] self._requires_grad=False" # t2 = prims.full((4096, 32, 128), 0, device=devices.Device("cuda:0"), dtype=dtypes.bfloat16) # t2: "cuda:0 bf16[4096, 32, 128] self._requires_grad=False" # <eval_with_key>.83:8: _set_grad_enabled_1 = torch._C._set_grad_enabled(True); _set_grad_enabled_1 = None (t0,) = prims.update_aliases((t2,)) t4 = ltorch.setitem_(t0, t_1, t_0) # t4: "cuda:0 bf16[4096, 32, 128] self._requires_grad=False" # t3 = ltorch.setitem(t0, t_1, t_0) # t3: "cuda:0 bf16[4096, 32, 128] self._requires_grad=False" # t3 = prims.copy_with_setitem(t0, t_1, t_0) # t3: "cuda:0 bf16[4096, 32, 128] self._requires_grad=False" # t4 = prims.copy_(t3, t0, grad_enabled=False) # t4: "cuda:0 bf16[4096, 32, 128] self._requires_grad=False" ltorch._set_grad_enabled_with_warning(True) return (t4, [t_1]) ``` seems to return `None` inside https://github.com/Lightning-AI/lightning-thunder/blob/a7952e29bae925581d3bf7069d64386f4d888e00/thunder/core/jit_ext.py#L998-L1007 of `_general_jit_torch_ops_higher_order_autograd_function_apply` and the cause seems to be that `ltorch.setitem_` has been a void function. Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
1 parent a7952e2 commit 71a3b20

1 file changed

Lines changed: 1 addition & 1 deletion

File tree

thunder/torch/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1141,7 +1141,7 @@ def setitem(inp, idx, val):
11411141

11421142
@torchsymbol(torch.Tensor.__setitem__, id="setitem_", is_method=True, tags=(prims.OpTags.IN_PLACE,))
11431143
def setitem_(inp, idx, val):
1144-
_copy_(inp, setitem(inp, idx, val))
1144+
return _copy_(inp, setitem(inp, idx, val))
11451145

11461146

11471147
@torchsymbol(torch.Tensor.__getitem__, id="torch.Tensor.__getitem__", method_name="getitem")

0 commit comments

Comments
 (0)