quaterion_models.heads.gated_head module

class GatedHead(input_embedding_size: int, dropout: float = 0.0)

Bases: EncoderHead

Disables or amplifies some components of input embedding.

This layer has minimal amount of trainable parameters and is suitable for even small training sets.

reset_parameters() → None

transform(input_vectors: Tensor) → Tensor

Parameters:

input_vectors – shape: (batch_size, vector_size)

Returns:

Tensor – (batch_size, vector_size)

property output_size: int

training: bool