diff --git a/1.1/HalPolicy.cpp b/1.1/HalPolicy.cpp
index 1b1c06e..d241290 100644
--- a/1.1/HalPolicy.cpp
+++ b/1.1/HalPolicy.cpp
@@ -300,19 +300,19 @@
     if(!axisOperand)
     {
         axis.assign(dimensionSequence,
-                    dimensionSequence+inputInfo.GetNumDimensions());
+                    dimensionSequence+rank);
     }
     else
     {
         GetTensorInt32Values(*axisOperand, axis, model, data);
     }
 
-    std::vector<uint32_t> outputDims;
-    for (auto& i : axis)
-    {
-        auto currentDimension = inputInfo.GetShape()[i];
-        bool skipSqueeze = (std::find(axis.begin(), axis.end(), i) == axis.end());
 
+    std::vector<uint32_t> outputDims;
+    for(unsigned int i = 0; i < rank; i++)
+    {
+        bool skipSqueeze = (std::find(axis.begin(), axis.end(), i) == axis.end());
+        auto currentDimension = inputInfo.GetShape()[i];
         if (skipSqueeze || currentDimension != 1)
         {
             outputDims.push_back(currentDimension);