| 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() |
no outgoing calls