| 130 | # with_kwargs is only supported in PyTorch 2.0 |
| 131 | # use this Catcher hack for now |
| 132 | class Catcher(nn.Module): |
| 133 | def __init__(self, module): |
| 134 | super().__init__() |
| 135 | self.module = module |
| 136 | |
| 137 | def forward(self, inp, **kwargs): |
| 138 | inps.append(inp) |
| 139 | layer_kwargs.update(kwargs) |
| 140 | raise ValueError # early exit to break later inference |
| 141 | |
| 142 | # patch layer 0 to catch input and kwargs |
| 143 | layers[0] = Catcher(layers[0]) |