Determine layout

This commit is contained in:
Liubov Batanina 2021-03-04 13:05:01 +03:00
parent 0689c70dba
commit 94533e12eb

View File

@ -769,8 +769,14 @@ static InferenceEngine::Layout estimateLayout(const Mat& m)
{ {
if (m.dims == 4) if (m.dims == 4)
return InferenceEngine::Layout::NCHW; return InferenceEngine::Layout::NCHW;
else if (m.dims == 3)
return InferenceEngine::Layout::CHW;
else if (m.dims == 2) else if (m.dims == 2)
return InferenceEngine::Layout::NC; return InferenceEngine::Layout::NC;
else if (m.dims == 1)
return InferenceEngine::Layout::C;
else if (m.dims == 5)
return InferenceEngine::Layout::NCDHW;
else else
return InferenceEngine::Layout::ANY; return InferenceEngine::Layout::ANY;
} }