Merge pull request BVLC#2 from TUNG/fix-resnet50-convergence

fix the problem of convergence for resnet50
tungld · Mar 24, 2017 · dbae7ab · dbae7ab
2 parents f62effc + 8cf0da3
commit dbae7ab
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/caffe/layers/scale_layer.cu b/src/caffe/layers/scale_layer.cu
@@ -128,6 +128,7 @@ void ScaleLayer<Dtype>::Backward_gpu(const vector<Blob<Dtype>*>& top,
         <<<CAFFE_GET_BLOCKS(count), CAFFE_CUDA_NUM_THREADS>>>(
         count, top_diff, scale_data, scale_dim_, inner_dim_, bottom_diff);
   }
+  CUDA_CHECK(cudaStreamSynchronize(cudaStreamDefault));
 }
 
 INSTANTIATE_LAYER_GPU_FUNCS(ScaleLayer);