diff --git a/TensorFlowSharp/Tensorflow.cs b/TensorFlowSharp/Tensorflow.cs
index b09dc8f8..a3ed344b 100644
--- a/TensorFlowSharp/Tensorflow.cs
+++ b/TensorFlowSharp/Tensorflow.cs
@@ -1917,7 +1917,8 @@ public TFOperationDesc SetAttrShape (string attrName, TFShape [] shapeList)
 					IntPtr array = Marshal.AllocHGlobal (sizeof (long) * shapeList [i].dims.Length);
 					Marshal.Copy (shapeList [i].dims, 0, array, shapeList [i].dims.Length);
 					Marshal.WriteIntPtr (unmanaged, ofs, array);
-					ofs += sizeof (IntPtr);
+				    num_dims[i] = shapeList[i].NumDimensions;
+                    ofs += sizeof (IntPtr);
 				}
 				TF_SetAttrShapeList (handle, attrName, unmanaged, num_dims, num_shapes);
 				ofs = 0;