提交 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 {
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(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) {
lod_tensor.mutable_data<float>(place);
lod_tensor.set_lod(src_lod);
CHECK_EQ(lod_tensor.lod_element(0, 2), 4UL);
CHECK_EQ(lod_tensor.lod_element(0, 4), 8UL);
CHECK_EQ(lod_tensor.lod_element(0, 2).first, 4UL);
CHECK_EQ(lod_tensor.lod_element(0, 4).first, 8UL);
auto lod = lod_tensor.lod();
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册