在提交新问题之前，请确保您已经在社区中搜索过相关问题，并使用了社区中提供的资源/工具后，仍未找到满意的解决方式。

环境信息

例如：
- 操作系统
- 昇腾硬件信息
- CANN软件版本
- 安装的对应软件版本

使用场景及问题

对于用户自定义的Tensor无法采集，预期是采集结果正常，但是实际采集不到module的输入和输出。

如下MyTensor是用户自己定义的Tensor，其实将torch.Tensor包了一层：

class MyTensor:
    def __init__(self, tensor: Optional[torch.Tensor]):
        self.actual_tensor: Optional[torch.Tensor] = tensor

    def __repr__(self):
        return f"MyTensor({self.actual_tensor})"

    def __str__(self):
        return str(self.actual_tensor)

    def __getattr__(self, name):
        if self.actual_tensor is None:
            return super().__getattr__(name)
        else:
            original_attr = getattr(self.actual_tensor, name)
            if isinstance(original_attr, Callable):
                def wrapper(*args, **kwargs):
                    args, kwargs = tree_map_only(MyTensor, lambda x: x.actual_tensor, (args, kwargs))
                    results = original_attr(*args, **kwargs)
                    return tree_map_only(torch.Tensor, lambda x: MyTensor(x), results)
                return wrapper
            elif original_attr is not None:
                return MyTensor(original_attr) if isinstance(original_attr, torch.Tensor) else original_attr
            else:
                return super().__getattr__(name)
    
def mytensor(
    data: Any,
    requires_grad: bool = False,
    *args,
    **kwargs,
）-> MyTensor:
    return MyTensor(torch.tensor(data, requires_grad=requires_grad, *args, **kwargs))

触发问题的简单示例代码:

from msprobe.pytorch import PrecisionDebugger

debugger = PrecisionDebugger(config_path="./config.json")

class SimpleAddModel(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, input_ids, attention_mask=None, *args, **kwargs):
        # 最简单的逻辑：输入 + 输入
        x = input_ids.actual_tensor
        out = x + x  # 加法操作
        
        return MyTensor(out)

if __name__ == "__main__":
    model = SimpleAddModel()
    
    input_ids = mytensor([[1, 2], [3, 4]])
    attention_mask = mytensor([[1, 1], [1, 1]])
    
    debugger.start(model=self.model)
    model_output = self.model(input_ids=input_ids, attention_mask=attention_mask, *args, **kwargs) # 由于返回的model_output为MyTensor的对象，此时采集不到module数据
    debugger.stop()
    debugger.step()

欢迎加入社区，感谢您对社区的贡献 🎉!