Commit fd2234e
committed
Return the updated
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>inp from seteitem_
1 parent b8d644d commit fd2234e
1 file changed
Lines changed: 1 addition & 1 deletion
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
1141 | 1141 | | |
1142 | 1142 | | |
1143 | 1143 | | |
1144 | | - | |
| 1144 | + | |
1145 | 1145 | | |
1146 | 1146 | | |
1147 | 1147 | | |
| |||
0 commit comments