Fix NERPredictor for zero initialization

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/16931

Reviewed By: dragonxlwang

Differential Revision: D14016749

fbshipit-source-id: b5512c52cef77651bdba1e31f588ea649daacdd9

diff --git a/caffe2/operators/pack_segments.cc b/caffe2/operators/pack_segments.cc
index 8e82ed0..edbaa16 100644 (file)
--- a/caffe2/operators/pack_segments.cc
+++ b/caffe2/operators/pack_segments.cc
@@ -67,11 +67,6 @@ bool PackSegmentsOp<CPUContext>::DoRunWithType2() {
     presence_mask_data = presence_mask->template mutable_data<bool>();
   }
 
-  if (!data.size(0)) {
-    // Return empty output (with the proper shape)
-    return true;
-  }
-
   // Do padding
   if (output->template IsType<float>()) {
     math::Set<float, CPUContext>(
@@ -80,6 +75,11 @@ bool PackSegmentsOp<CPUContext>::DoRunWithType2() {
         output->template mutable_data<float>(),
         &context_);
   }
+  if (!data.size(0)) {
+    // Return empty output (with the proper shape)
+    return true;
+  }
+
   if (return_presence_mask_) {
     memset(presence_mask_data, (int)false, presence_mask->numel());
   }