torch_head_dimension

probly.utils.torch.torch_head_dimension(head: Module, name: Literal['in_features', 'out_features']) → int[source]

Read an integer feature dimension from a registered classification head.

Custom heads may use any module class with the requested attribute. Integer-like values implementing __index__ are supported, including NumPy integers.

Parameters:

  • head – Classification head selected by a traversal registration.

  • name – Feature dimension needed by the consumer.

Returns:

The requested feature dimension.

Raises:

TypeError – If the head does not expose the requested integer dimension.