[WIP] [pfto] Add ppe.map test case

pfnet · Jul 4, 2023 · 30acedc · 30acedc
1 parent ef49153
commit 30acedc
Showing 1 changed file with 15 additions and 0 deletions.
diff --git a/tests/pytorch_pfn_extras_tests/onnx_tests/test_export.py b/tests/pytorch_pfn_extras_tests/onnx_tests/test_export.py
@@ -389,3 +389,18 @@ def forward(self, x):
     assert len(model.graph.input) == 1
     model = run_model_test(Model(), (torch.rand((1,)),), keep_initializers_as_inputs=True)
     assert len(model.graph.input) == (2 if persistent else 1)
+
+
+def test_ppe_map():
+    torch.manual_seed(100)
+
+    class Net(torch.nn.Module):
+        def __init__(self):
+            super(Net, self).__init__()
+            self.conv = torch.nn.Conv2d(1, 1, 3)
+
+        def forward(self, x):
+            y = self.conv(x)
+            return list(ppe.map(lambda u: u + 1, y))[0]
+
+    run_model_test(Net(), (torch.rand(1, 1, 112, 112),), rtol=1e-03)