使用pybind11绑定C++返回torch::Tensor时部分尺寸触发段错误

你的核心问题出在**torch::from_blob的内存所有权逻辑**以及data_array的内存分配方式上：

• torch::from_blob创建的张量仅作为现有内存的视图（View），不会主动接管输入内存的所有权。如果data_array是在processing函数内部用栈分配的局部数组（比如uint16_t data_array[XXX];），函数返回后栈内存会被操作系统回收，此时Python端拿到的张量指向的是已释放的野内存。
• 不同num_neurons对应的总数据量不同：总元素数 = batch_size × num_neurons × data_len × data_wid。当总数据量超过Ubuntu进程默认栈大小（通常为8MB）时，会触发栈内存溢出或后续内存被覆盖，访问时就会出现段错误；数据量较小时栈内存暂时未被覆盖，因此能正常访问——这就是仅特定num_neurons取值会出错的原因。

另外你代码里存在一个隐性错误：torch::kUInt8数据类型和data_array的uint16_t类型不匹配，这会导致数据解析异常，必须修正为一致类型。

方案1：改用堆分配+自定义销毁器 #

将data_array改为堆分配，同时给from_blob指定销毁器，让张量销毁时自动释放堆内存：

torch::Tensor processing(py::array np_data, py::array dim){
    // 解析目标形状并计算总元素数
    std::vector<int64_t> target_shape;
    for (auto elem : dim) {
        target_shape.push_back(static_cast<int64_t>(elem));
    }
    int64_t total_elements = 1;
    for (auto s : target_shape) total_elements *= s;

    // 堆分配数据数组
    uint16_t* data_array = new uint16_t[total_elements];

    // 将numpy数组转换到data_array的逻辑
    // ...

    // 指定匹配的数据类型，同时添加销毁器
    auto options = torch::TensorOptions().dtype(torch::kUInt16);
    torch::Tensor output = torch::from_blob(data_array, {total_elements}, options, [](void* ptr) {
        delete[] static_cast<uint16_t*>(ptr);
    });
    auto finalOutput = torch::reshape(output, target_shape);
    return finalOutput;
}

方案2：直接用Torch API创建张量（推荐） #

避免手动管理内存，直接创建Torch张量并拷贝数据：

torch::Tensor processing(py::array np_data, py::array dim){
    // 解析目标形状
    std::vector<int64_t> target_shape;
    for (auto elem : dim) {
        target_shape.push_back(static_cast<int64_t>(elem));
    }

    // 创建空的目标张量
    auto options = torch::TensorOptions().dtype(torch::kUInt16);
    torch::Tensor finalOutput = torch::empty(target_shape, options);

    // 获取张量数据指针，直接拷贝数据
    uint16_t* tensor_data = finalOutput.data_ptr<uint16_t>();
    // 将numpy数组的数据拷贝到tensor_data的逻辑
    // ...

    return finalOutput;
}

额外修正点 #

你的Python代码中output_shape = np.ndarray([batch_size, num_neurons, data_len, data_wid])是错误的，np.ndarray是底层构造函数，参数格式不对会生成空数组，导致C++端解析维度错误，应改为：

output_shape = np.array([batch_size, num_neurons, data_len, data_wid])

内容的提问来源于stack exchange，提问作者inventi