merge tests for topk

BowenBao · BowenBao · commit e76ae8e4ad6a · 2019-05-13T10:43:47.000-07:00
diff --git a/test/onnx/test_pytorch_onnx_caffe2.py b/test/onnx/test_pytorch_onnx_caffe2.py
@@ -1356,23 +1356,6 @@ def forward(self, input):
         caffe2_out = prepared.run(inputs=[x.cpu().numpy()])
         self.assertEqual(caffe2_out[0].shape, x.shape)
 
-    def test_topk(self):
-        class TopKModel(torch.nn.Module):
-            def forward(self, input):
-                return torch.topk(input, 3, dim=0)
-
-        x = torch.randn(4, 3, requires_grad=True)
-        self.run_model_test(TopKModel(), train=False, input=x, batch_size=BATCH_SIZE)
-
-    def test_topk_script(self):
-        class TopKModel(torch.jit.ScriptModule):
-            @torch.jit.script_method
-            def forward(self, input):
-                return torch.topk(input, 3, dim=0)
-
-        x = torch.randn(4, 3, requires_grad=True)
-        self.run_model_test(TopKModel(), train=False, input=(x,), batch_size=BATCH_SIZE, example_outputs=torch.topk(x, 3, dim=0))
-
     def test_traced_ints(self):
         A = 4
         H = 10
@@ -1600,10 +1583,19 @@ def test_topk(self):
         class TopKModel(torch.nn.Module):
             def forward(self, input):
                 return torch.topk(input, 3)
-        model = TopKModel()
+
         x = torch.arange(1., 6.)
         self.run_model_test(TopKModel(), train=False, input=x, batch_size=BATCH_SIZE)
 
+    def test_topk_script(self):
+        class TopKModel(torch.jit.ScriptModule):
+            @torch.jit.script_method
+            def forward(self, input):
+                return torch.topk(input, 3, dim=0)
+
+        x = torch.randn(4, 3, requires_grad=True)
+        self.run_model_test(TopKModel(), train=False, input=(x,), batch_size=BATCH_SIZE, example_outputs=torch.topk(x, 3, dim=0))
+
     def test_floor(self):
         class FloorModel(torch.nn.Module):
             def forward(self, input):
