Commit 4d59204f authored by charlie's avatar charlie
Browse files

Simplify compare to 1 dyn_dim

parent ed2c235e
...@@ -59,9 +59,9 @@ struct squeeze ...@@ -59,9 +59,9 @@ struct squeeze
auto input_shape = inputs[0]; auto input_shape = inputs[0];
if(input_shape.dynamic()) if(input_shape.dynamic())
{ {
std::vector<shape::dynamic_dimension> one_dyn_dims{{1, 1, 0}, {1, 1, 1}}; shape::dynamic_dimension one_dyn_dim{1, 1, 0};
if(std::any_of(axes.begin(), axes.end(), [&](auto axis) { if(std::any_of(axes.begin(), axes.end(), [&](auto axis) {
return not contains(one_dyn_dims, input_shape.dyn_dims()[axis]); return input_shape.dyn_dims()[axis] != one_dyn_dim;
})) }))
{ {
MIGRAPHX_THROW( MIGRAPHX_THROW(
...@@ -73,7 +73,7 @@ struct squeeze ...@@ -73,7 +73,7 @@ struct squeeze
for(auto i : range(input_shape.ndim())) for(auto i : range(input_shape.ndim()))
{ {
auto dd = input_shape.dyn_dims()[i]; auto dd = input_shape.dyn_dims()[i];
if(not contains(one_dyn_dims, dd)) if(dd != one_dyn_dim)
{ {
dyn_dims.push_back(dd); dyn_dims.push_back(dd);
} }
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment