提交 c8c4ba0f 编写于 作者: W whs 提交者: GitHub

Merge pull request #4918 from wanghaoshuang/fix_lod

Make LoDTensor::lod_element return start and end offset of sequence
...@@ -74,12 +74,12 @@ class LoDTensor : public Tensor { ...@@ -74,12 +74,12 @@ class LoDTensor : public Tensor {
LoD lod() const { return lod_; } LoD lod() const { return lod_; }
/* /*
* Get a element from LoD. * Get the start offset and end offset of an element from LoD.
*/ */
size_t lod_element(size_t level, size_t elem) const { std::pair<size_t, size_t> lod_element(size_t level, size_t elem) const {
PADDLE_ENFORCE_LT(level, NumLevels()); PADDLE_ENFORCE_LT(level, NumLevels());
PADDLE_ENFORCE_LT(elem, NumElements(level)); PADDLE_ENFORCE_LT(elem, NumElements(level));
return (lod_)[level][elem]; return std::make_pair((lod_)[level][elem], (lod_)[level][elem + 1]);
} }
/* /*
......
...@@ -36,8 +36,8 @@ TEST(LoDTensor, LoDInGPU) { ...@@ -36,8 +36,8 @@ TEST(LoDTensor, LoDInGPU) {
lod_tensor.mutable_data<float>(place); lod_tensor.mutable_data<float>(place);
lod_tensor.set_lod(src_lod); lod_tensor.set_lod(src_lod);
CHECK_EQ(lod_tensor.lod_element(0, 2), 4UL); CHECK_EQ(lod_tensor.lod_element(0, 2).first, 4UL);
CHECK_EQ(lod_tensor.lod_element(0, 4), 8UL); CHECK_EQ(lod_tensor.lod_element(0, 4).first, 8UL);
auto lod = lod_tensor.lod(); auto lod = lod_tensor.lod();
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册