commit c33c9c1ba452dbfd50ac87a9915ac14fb45decca
author Frank Jiang <ffjiang@fb.com> Wed Sep 27 18:58:03 2017 -0700
committer Facebook Github Bot <facebook-github-bot@users.noreply.github.com> Wed Sep 27 19:05:56 2017 -0700
tree 2e4a444b2455df765a2cf999e99e1554c6f05ee5
parent 2108d1c250690a76ff17dbe31cebfe9940876ebc

Fixed size_to_dim enforce

Summary: Fixed Caffe2Enforce in size_to_dim() so that it works even if k is same as the number of dimensions in the tensor.

Reviewed By: salexspb

Differential Revision: D5893264

fbshipit-source-id: 525ea263f5e21e197c7010e1c66501355b8027c8

caffe2/core/tensor.h

diff --git a/caffe2/core/tensor.h b/caffe2/core/tensor.h
index c31676d..7df4837 100644
--- a/caffe2/core/tensor.h
+++ b/caffe2/core/tensor.h
@@ -47,7 +47,7 @@
 
 // Product of all dims up to
 inline TIndex size_to_dim_(int k, vector<TIndex> dims) {
-  CAFFE_ENFORCE(k < dims.size());
+  CAFFE_ENFORCE(k <= dims.size());
   TIndex r = 1;
   for (int i = 0; i < k; ++i) {
     r *= dims[i];