add test for test

Author: xadupre

_unittests/ut_reference/test_onnxruntime_evaluator.py

@@ -59,6 +59,43 @@ def forward(self, x):
         got = orte.run(None, {name: x.numpy()})[0]
         self.assertEqualArray(expected, got)
 
+    def test_ort_eval_cond(self):
+        import torch
+
+        class TwoInputs(torch.nn.Module):
+            def forward(self, x, y):
+                def true_fn(x, y):
+                    return torch.sin(x), torch.cos(x) + y
+
+                def false_fn(x, y):
+                    return torch.cos(x), torch.sin(x) + y
+
+                return torch.cond(x.sum() > 0, true_fn, false_fn, [x, y])
+
+        x, y = torch.rand(5, 3), torch.rand(5, 3)
+        model = TwoInputs()
+        onx = to_onnx(model, (x, y), inline=False)
+        self.assertEqual(len(onx.functions), 2)
+
+        # ExtendedReferenceEvaluator
+        ref = ExtendedReferenceEvaluator(onx)
+        for _x in (x, -x):
+            expected = model(_x, y)
+            got = ref.run(None, {"x": _x.detach().numpy(), "y": y.detach().numpy()})
+            self.assertEqual(len(expected), len(got))
+            for e, g in zip(expected, got):
+                self.assertEqualArray(e, g, atol=1e-5)
+
+        # OnnxruntimeEvaluator
+        ref = OnnxruntimeEvaluator(onx)
+
+        for _x in (x, -x):
+            expected = model(_x, y)
+            got = ref.run(None, {"x": _x.detach().numpy(), "y": y.detach().numpy()})
+            self.assertEqual(len(expected), len(got))
+            for e, g in zip(expected, got):
+                self.assertEqualArray(e, g, atol=1e-5)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)