fix tensor_info

Neargye · Dec 21, 2018 · 99295fe · 99295fe
1 parent 1e63b04
commit 99295fe
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/src/tensor_info.cpp b/src/tensor_info.cpp
@@ -101,6 +101,13 @@ void PrintOutputs(TF_Graph* graph, TF_Operation* op) {
       continue;
     }
 
+    std::cout << " dims: " << num_dims;
+
+    if (num_dims <= 0) {
+      std::cout << " []" << std::endl;;
+      continue;
+    }
+
     std::vector<std::int64_t> dims(num_dims);
 
     std::cout << "Output: " << i << " type: " << TFDataTypeToString(type);
@@ -111,7 +118,7 @@ void PrintOutputs(TF_Graph* graph, TF_Operation* op) {
       continue;
     }
 
-    std::cout << " dims: " << num_dims << " [";
+    std::cout << " [";
     for (int d = 0; d < num_dims; ++d) {
       std::cout << dims[d];
       if (d < num_dims - 1) {