more fix

SPRCSY · Sep 18, 2015 · 3c42e9f · 3c42e9f
1 parent 4a8c3a3
commit 3c42e9f
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 4 deletions.
diff --git a/mshadow-ps/ps_local-inl.h b/mshadow-ps/ps_local-inl.h
@@ -630,9 +630,9 @@ class LocalModel : public ISharedModel<xpu, DType> {
         PushEntry &e = push_map.GetRef(tsk.key);
         CHECK_EQ(e.data[0][0].shape_, tsk.data.shape_)
           << "Tensor with same key must share same shape "
-          << e.data[0][0].shape_.ToString()
+          << e.data[0][0].shape_
           << " vs "
-          << tsk.data.shape_.ToString();
+          << tsk.data.shape_;
         CHECK_EQ(!e.copied[wid], true) << "data inconsistency";
         // start copy
         SetDevice<xpu>(tsk.devid);

diff --git a/mshadow/tensor.h b/mshadow/tensor.h
@@ -12,6 +12,7 @@
 #ifndef MSHADOW_TENSOR_H_
 #define MSHADOW_TENSOR_H_
 #include <string>
+#include <iostream>
 #include "./base.h"
 #include "./expression.h"
 
@@ -30,6 +31,18 @@ struct gpu {
   /*! \brief device flag number, identifies this device */
   static const int kDevMask = 1 << 1;
 };
+template<int ndim>
+struct Shape;
+
+/*!
+ * \brief allow string printing of the shape
+ * \param os the output stream
+ * \param shape the shape
+ * \return the ostream
+ */
+template<int ndim>
+inline std::ostream &operator<<(std::ostream &os, const Shape<ndim> &shape); // NOLINT(*)
+
 /*!
  * \brief shape of a tensor
  *       IMPORTANT NOTE: this shape is different from numpy.shape
@@ -153,8 +166,8 @@ v   * \return subshape
     }
     return s;
   }
-  // print function
-  friend std::ostream &operator<<(std::ostream &os, const Shape<dimension> &shape); // NOLINT(*)
+  template<int dim>
+  friend std::ostream &operator<<(std::ostream &os, const Shape<dim> &shape); // NOLINT(*)
 };  // Shape
 //------------------------------------------------
 // useful construction functions to generate shape