1# mypy: allow-untyped-defs 2import torch 3 4class Dictionary(torch.nn.Module): 5 """ 6 Dictionary structures are inlined and flattened along tracing. 7 """ 8 9 def forward(self, x, y): 10 elements = {} 11 elements["x2"] = x * x 12 y = y * elements["x2"] 13 return {"y": y} 14 15example_args = (torch.randn(3, 2), torch.tensor(4)) 16tags = {"python.data-structure"} 17model = Dictionary() 18