RuntimeError in Graph Mode with flow.unique Operator

[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 calling flow.unique with return_inverse=True within the forward method.

Code to reproduce bug

import oneflow as flow
import oneflow.nn as nn

class TestModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.relu = nn.ReLU()

    def forward(self, x):
        _x, _i = flow.unique(x, sorted=True, return_inverse=True)
        _x = self.relu(_x)
        return flow.mean(_x), _i

# Test data
x = flow.randn(32, 64)

model = TestModel()

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

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

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

Error Logs

[ERROR](GRAPH:TestGraph_0:TestGraph) building graph got error.
Traceback (most recent call last):
  File "test_script.py", line 39, in <module>
    print("Graph mode result:", g(x))
  ... (include the complete traceback from your error)
  File "test_script.py", line 13, in forward
    _x, _i = flow.unique(x, sorted=True, return_inverse=True)
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