Not changing dimension size for expand when target size is -1

1 files changed, 10 insertions, 0 deletions

diff --git a/lib/TH/THStorage.c b/lib/TH/THStorage.c
index f6b63f4..e57181c 100644
--- a/lib/TH/THStorage.c
+++ b/lib/TH/THStorage.c
@@ -132,6 +132,16 @@ int THLongStorage_inferExpandGeometry(long *tensorSizes, long *tensorStrides, lo
     long stride = (dim >= 0) ?
         tensorStrides[dim] : expandedSizesCalc[i + 1] * expandedStridesCalc[i+1];
     long targetSize = THLongStorage_data(sizes)[i];
+    if (targetSize == -1) {
+      if (dim < 0) {
+        THFree(expandedSizesCalc);
+        THFree(expandedStridesCalc);
+        snprintf(error_buffer, buffer_len, "The expanded size of the tensor (%ld) isn't allowed in a leading, non-existing dimension %ld.", targetSize, i);
+        return -1;
+      } else {
+        targetSize = size;
+      }
+    }
     if (size != targetSize) {
       if (size == 1) {
         size = targetSize;