1# mypy: allow-untyped-defs 2import torch 3 4class StaticForLoop(torch.nn.Module): 5 """ 6 A for loop with constant number of iterations should be unrolled in the exported graph. 7 """ 8 9 def forward(self, x): 10 ret = [] 11 for i in range(10): # constant 12 ret.append(i + x) 13 return ret 14 15example_args = (torch.randn(3, 2),) 16tags = {"python.control-flow"} 17model = StaticForLoop() 18