RuntimeError in Graph Mode with Custom Tensor Subclasses

[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 creating instances of custom Tensor subclasses (MyTensor) within the forward method during graph compilation.

Code to reproduce bug

import oneflow as flow
import oneflow.nn as nn

class MyTensor(flow.Tensor):
    pass

class TestModel(flow.nn.Module):
    def forward(self, t: flow.Tensor):
        my_tensor = MyTensor(flow.randn(128, 64))
        return my_tensor * flow.sigmoid(t)

# Test data
t = flow.randn(32, 128, 64)
model = TestModel()

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

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

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

Error Logs

[ERROR](GRAPH:TestGraph_0:TestGraph) building graph got error.
Traceback (most recent call last):
  File "test_script.py", line 37, in <module>
    print("Graph mode result:", g(t))
  ...
  File "test_script.py", line 12, in forward
    my_tensor = MyTensor(flow.randn(128, 64))
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