Method test_ops

caffe2/python/net_builder_test.py:83–104

(self)

class TestNetBuilder(unittest.TestCase):
    def test_ops(self):
        with NetBuilder() as nb:
            y = _test_loop()
            z, w, a, b = _test_outer()
            p = _test_if(ops.Const(75))
            q = _test_if(ops.Const(25))
        plan = Plan('name')
        plan.AddStep(to_execution_step(nb))
        ws = workspace.C.Workspace()
        ws.run(plan)
        expected_results = [
            (y, 5),
            (z, False),
            (w, True),
            (a, False),
            (b, False),
            (p, 2),
            (q, 3),
        ]
        for b, expected in expected_results:
            actual = ws.blobs[str(b)].fetch()
            self.assertEqual(actual, expected)

    def _expected_loop(self):
        total = 0