find_layer

probly.traverse_nn.utils.find_layer(model: object, layer_types: type | tuple[type, ...]) → Any[source]

Return the first layer of the given type(s) contained in a model.

Convenience wrapper around find_layers() for the common case of a model holding a single layer of interest (e.g. a swapped-in last layer).

Parameters:

  • model – The model to search.

  • layer_types – A layer type or tuple of layer types to match.

Returns:

The first matching layer in forward DFS order.

Raises:

ValueError – If the model contains no layer of the given type(s).