| 104 | # with_kwargs is only supported in PyTorch 2.0 |
| 105 | # use this Catcher hack for now |
| 106 | class Catcher(nn.Module): |
| 107 | def __init__(self, module): |
| 108 | super().__init__() |
| 109 | self.module = module |
| 110 | |
| 111 | def forward(self, inp, **kwargs): |
| 112 | inps.append(inp) |
| 113 | layer_kwargs.update(kwargs) |
| 114 | raise ValueError # early exit to break later inference |
| 115 | |
| 116 | # patch layer 0 to catch input and kwargs |
| 117 | layers[0] = Catcher(layers[0]) |