fixup! SoftmaxLossLayer generalized like SoftmaxLayer

BVLC · Mar 3, 2015 · 025dd23 · 025dd23
1 parent 85518e4
commit 025dd23
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/caffe/layers/softmax_loss_layer.cpp b/src/caffe/layers/softmax_loss_layer.cpp
@@ -38,6 +38,11 @@ void SoftmaxWithLossLayer<Dtype>::Reshape(
   softmax_axis_ = this->layer_param_.softmax_param().axis();
   outer_num_ = bottom[0]->count(0, softmax_axis_);
   inner_num_ = bottom[0]->count(softmax_axis_ + 1);
+  CHECK_EQ(outer_num_ * inner_num_, bottom[1]->count())
+     << "Number of labels must match number of predictions; "
+     << "e.g., if softmax axis == 1 and prediction shape is (N, C, H, W), "
+     << "label count (number of labels) must be N*H*W, "
+     << "with integer values in {0, 1, ..., C-1}.";
   if (top.size() >= 2) {
     // softmax output
     top[1]->ReshapeLike(*bottom[0]);