| 107 | |
| 108 | |
| 109 | class MultiEntryPointStatefulModel(torch.nn.Module): |
| 110 | def __init__(self) -> None: |
| 111 | super().__init__() |
| 112 | self.register_buffer("state", torch.zeros(2, 2)) |
| 113 | |
| 114 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 115 | return self.state.add_(x).view(-1) * 2 |
| 116 | |
| 117 | def set_state(self, state: torch.Tensor) -> None: |
| 118 | self.state.copy_(state) |
| 119 | |
| 120 | def get_state(self) -> torch.Tensor: |
| 121 | return self.state |
| 122 | |
| 123 | def get_example_inputs(self) -> Tuple[torch.Tensor, ...]: |
| 124 | return (torch.ones(1),) |
| 125 | |
| 126 | |
| 127 | class ModelWithDifferentTensorSizes(torch.nn.Module): |
no outgoing calls