1# mypy: allow-untyped-defs 2import torch 3 4class DynamicShapeConstructor(torch.nn.Module): 5 """ 6 Tensor constructors should be captured with dynamic shape inputs rather 7 than being baked in with static shape. 8 """ 9 10 def forward(self, x): 11 return torch.zeros(x.shape[0] * 2) 12 13example_args = (torch.randn(3, 2),) 14tags = {"torch.dynamic-shape"} 15model = DynamicShapeConstructor() 16