[MSLITE][Develop] fix bug of tensor_rt

This commit is contained in:
yangruoqi713 2021-07-12 09:58:15 +08:00
parent 6c7982f18a
commit 492940d6e2
1 changed files with 1 additions and 1 deletions

View File

@ -33,7 +33,7 @@ int GatherTensorRT::IsSupport(const schema::Primitive *primitive, const std::vec
return RET_ERROR;
}
if (in_tensors[2].ElementNum() == 1) {
axis_ = static_cast<int *>(in_tensors[2].MutableData())[0];
axis_ = static_cast<const int *>(in_tensors[2].Data().get())[0];
} else {
MS_LOG(ERROR) << "TensorRT axis is attribute.";
return RET_ERROR;