RuntimeError in Graph Mode with Tensor Indexing Assignment in Forward Pass

[LiSsHhUuAaIi]

Summary

I encountered a RuntimeError when running a model in OneFlow's Graph mode that works perfectly in eager mode. The error occurs specifically when performing tensor indexing assignment operations like x[i] = flow.tanh(x[i]) within the forward method.

Code to reproduce bug

import oneflow as flow
import oneflow.nn as nn
import oneflow.nn.functional as F

class TestModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(64, 128, bias=True)

    def forward(self, x: flow.Tensor, i: flow.Tensor) -> flow.Tensor:
        x = F.relu(self.linear(x))
        # This tensor indexing assignment fails in Graph mode
        x[i] = flow.tanh(x[i])
        return x

# Test data
x = flow.randn(32, 64)
i = flow.tensor([0, 2, 5, 8, 12])

model = TestModel()

# Eager mode works
print("Eager mode result shape:", model(x, i).shape)

# Graph mode fails
class TestGraph(nn.Graph):
    def __init__(self, model):
        super().__init__()
        self.model = model
        
    def build(self, x, i):
        return self.model(x, i)

g = TestGraph(model)
print("Graph mode result:", g(x, i))

Error Logs

[ERROR](GRAPH:TestGraph_0:TestGraph) building graph got error.
Traceback (most recent call last):
  File "test_script.py", line 35, in <module>
    print("Graph mode result:", g(x, i))
  ...
  File "test_script.py", line 20, in forward
    x[i] = flow.tanh(x[i])
RuntimeError: Error: RuntimeError : This is a oneflow bug, please submit an issue at 'https://github.com/Oneflow-Inc/oneflow/issues' including the log information of the error, the minimum reproduction code, and the system information.

System Information

• OS: Ubuntu 20.04.6 LTS
• OneFlow version: '1.0.0.dev20251101+cpu'
• Python version: Python 3.10.19