forked from mindspore-Ecosystem/mindspore
[MSLITE][Develop] fix bug of tensor_rt
This commit is contained in:
parent
6c7982f18a
commit
492940d6e2
|
@ -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;
|
||||
|
|
Loading…
Reference in New Issue