Skip to content
Snippets Groups Projects
Commit c1ecff21 authored by Grégoire Kubler's avatar Grégoire Kubler
Browse files

fix : output dims of squeeze were not properly created

parent 9ba16cb8
No related branches found
No related tags found
1 merge request!102feat/operator_squeeze_unsqueeze
Pipeline #53944 passed
......@@ -388,15 +388,14 @@ TEST_CASE("[core/operator] Squeeze(forward)", "[Squeeze][forward]") {
int i = 0;
std::vector<DimSize_t> dims_out;
dims_out.reserve(dims_in.size());
std::copy_if(dims_in.begin(), dims_in.end(), std::back_inserter(dims_out),
[&dims_to_squeeze, &i](DimSize_t dim) {
bool ok = dim != 1 ||
!std::binary_search(dims_to_squeeze.begin(),
dims_to_squeeze.end(), i);
i++; // incrementing counter since C++ has not enumerate
// fctn (until C++23)
return ok;
});
for (DimIdx_t i = 0; i < dims_in.size(); ++i) {
if (dims_in[i] == 1 &&
std::find(dims_to_squeeze.begin(), dims_to_squeeze.end(), i) !=
dims_to_squeeze.end()) {
continue;
}
dims_out.push_back(dims_in[i]);
}
CHECK(op->forwardDims());
CHECK(op->getOutput(0)->dims() == dims_out);
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment