Merge pull request BVLC#4559 from fyu/loss_reshape

Fix Label Shape Check in LossLayer::Reshape
Sensenzhl · Aug 7, 2016 · 375003a · 375003a
2 parents 6431477 + 61e0165
commit 375003a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/caffe/layers/loss_layer.cpp b/src/caffe/layers/loss_layer.cpp
@@ -16,8 +16,8 @@ void LossLayer<Dtype>::LayerSetUp(
 template <typename Dtype>
 void LossLayer<Dtype>::Reshape(
     const vector<Blob<Dtype>*>& bottom, const vector<Blob<Dtype>*>& top) {
-  CHECK_EQ(bottom[0]->num(), bottom[1]->num())
-      << "The data and label should have the same number.";
+  CHECK_EQ(bottom[0]->shape(0), bottom[1]->shape(0))
+      << "The data and label should have the same first dimension.";
   vector<int> loss_shape(0);  // Loss layers output a scalar; 0 axes.
   top[0]->Reshape(loss_shape);
 }