| 259 | |
| 260 | |
| 261 | class _LayerHook: |
| 262 | def __init__(self, layer, idx, storage): |
| 263 | self._layer, self._idx, self._storage = layer, idx, storage |
| 264 | |
| 265 | def __call__(self, *args, **kwargs): |
| 266 | out = self._layer(*args, **kwargs) |
| 267 | self._storage[self._idx] = out[0] if isinstance(out, tuple) else out |
| 268 | return out |
| 269 | |
| 270 | def __getattr__(self, name): |
| 271 | return getattr(self._layer, name) |
| 272 | |
| 273 | |
| 274 | def _get_layers(model): |