diff --git a/tools/extract_features.cpp b/tools/extract_features.cpp
index ddbce1075ed..8b1f9e8880e 100644
--- a/tools/extract_features.cpp
+++ b/tools/extract_features.cpp
@@ -146,9 +146,9 @@ int feature_extraction_pipeline(int argc, char** argv) {
       int dim_features = feature_blob->count() / batch_size;
       const Dtype* feature_blob_data;
       for (int n = 0; n < batch_size; ++n) {
-        datum.set_height(dim_features);
-        datum.set_width(1);
-        datum.set_channels(1);
+        datum.set_height(feature_blob->height());
+        datum.set_width(feature_blob->width());
+        datum.set_channels(feature_blob->channels());
         datum.clear_data();
         datum.clear_float_data();
         feature_blob_data = feature_blob->cpu_data() +