Fix UpsampleNearest op CPU impl batch handling (pytorch#13002)

Summary:
Pull Request resolved: pytorch#13002

Batch dim wasn't handled in the CPU impl (will fail for inputs with N > 1).
Fixing that here.

Differential Revision: D10515159

fbshipit-source-id: ee7e4f489d2d4de793f550b31db7c0e2ba3651e8

modules/detectron/upsample_nearest_op.h

@@ -50,7 +50,7 @@ class UpsampleNearestOp final : public Operator<Context> {
       d2 = Y->dim32(1);
       d3 = Y->dim32(2);
     } else {
-      d1 = Y->dim32(1);
+      d1 = Y->dim32(0) * Y->dim32(1);
       d2 = Y->dim32(2);
       d3 = Y->dim32(3);
     }

modules/detectron/upsample_nearest_op_test.py

@@ -0,0 +1,43 @@
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import unittest
+
+import caffe2.python.hypothesis_test_util as hu
+import hypothesis.strategies as st
+import numpy as np
+from caffe2.python import core, dyndep
+from hypothesis import given
+
+
+dyndep.InitOpsLibrary("@/caffe2/modules/detectron:detectron_ops")
+
+
+class TestUpsampleNearestOp(hu.HypothesisTestCase):
+    @given(
+        N=st.integers(1, 3),
+        H=st.integers(10, 300),
+        W=st.integers(10, 300),
+        scale=st.integers(1, 3),
+        **hu.gcs
+    )
+    def test_upsample_nearest_op(self, N, H, W, scale, gc, dc):
+        C = 32
+        X = np.random.randn(N, C, H, W).astype(np.float32)
+        op = core.CreateOperator("UpsampleNearest", ["X"], ["Y"], scale=scale)
+
+        def ref(X):
+            outH = H * scale
+            outW = W * scale
+            outH_idxs, outW_idxs = np.meshgrid(
+                np.arange(outH), np.arange(outW), indexing="ij"
+            )
+            inH_idxs = (outH_idxs / scale).astype(np.int32)
+            inW_idxs = (outW_idxs / scale).astype(np.int32)
+            Y = X[:, :, inH_idxs, inW_idxs]
+            return [Y]
+
+        self.assertReferenceChecks(device_option=gc, op=op, inputs=[X], reference=ref)
+
+
+if __name__ == "__main__":
+    unittest.main()