DataCategorical.get_ml_heads

DataCategorical.get_ml_heads(head_layer_widths: List[int], last_core_layer_width: int, activation: torch.nn.Module, **kwargs) → Tuple[torch.nn.Module, torch.nn.Module][source]

Returns a fully-connected head with the appropriate number of in / out channels for encoding / decoding this feature

Parameters:

  • head_layer_widths (List[int]) – List of integers where the length denotes number of layers and values the width of the layers in the head

  • last_core_layer_width (int) – Dimensionality of autoencoder output to be decoded by this feature head

  • activation (nn.Module) – Activation function to be used in this head

Returns:

Tuple[nn.Module, nn.Module] – Tuple containing the input and output head.