| 315 | |
| 316 | |
| 317 | def test_split_dump(): |
| 318 | class SimpleNet(Module): |
| 319 | def __init__(self, num_segments: int = 3): |
| 320 | super().__init__() |
| 321 | self.num_segments = num_segments |
| 322 | |
| 323 | def forward(self, x): |
| 324 | x = F.split(x, self.num_segments, axis=1) |
| 325 | return x |
| 326 | |
| 327 | model = SimpleNet() |
| 328 | model.eval() |
| 329 | data = tensor(np.random.random((1, 12, 224, 224))) |
| 330 | |
| 331 | @trace(symbolic=True, capture_as_const=True) |
| 332 | def fun(data, *, net): |
| 333 | return net(data) |
| 334 | |
| 335 | x = fun(data, net=model) |
| 336 | fun.dump(io.BytesIO(), arg_names=["data"]) |
| 337 | |
| 338 | |
| 339 | @pytest.mark.parametrize("trace_mode", [False, True]) |