flambe.nn.softmax

Module Contents

class flambe.nn.softmax.SoftmaxLayer(input_size: int, output_size: int, mlp_layers: int = 1, mlp_dropout: float = 0.0, mlp_hidden_activation: Optional[nn.Module] = None, take_log: bool = True)[source]

Bases: flambe.nn.module.Module

Implement an SoftmaxLayer module.

Can be used to form a classifier out of any encoder. Note: by default takes the log_softmax so that it can be fed to the NLLLoss module. You can disable this behavior through the take_log argument.

forward(self, data: torch.Tensor)[source]

Performs a forward pass through the network.

Parameters:data (torch.Tensor) – input to the model of shape (*, input_size)
Returns:output – output of the model of shape (*, output_size)
Return type:torch.Tensor