pub enum LayerOutput {
Embedding(Tensor),
EncoderWithAttention(HiddenLayer),
}
Expand description
Output of a BERT layer.
Variants§
Implementations§
source§impl LayerOutput
impl LayerOutput
sourcepub fn attention(&self) -> Option<&Tensor>
pub fn attention(&self) -> Option<&Tensor>
Get the layer attention.
Return a Some
value if the layer output is from an encoder layer,
or None
otherwise.
sourcepub fn embedding(&self) -> Option<&Tensor>
pub fn embedding(&self) -> Option<&Tensor>
Get the embedding.
Returns Some
if the layer output is an embedding or None
otherwise.
sourcepub fn map_output<F>(&self, f: F) -> Result<Self, TransformerError>where
F: Fn(&Tensor) -> Result<Tensor, TransformerError>,
pub fn map_output<F>(&self, f: F) -> Result<Self, TransformerError>where
F: Fn(&Tensor) -> Result<Tensor, TransformerError>,
Map the output representation of this layer.
sourcepub fn output_mut(&mut self) -> &mut Tensor
pub fn output_mut(&mut self) -> &mut Tensor
Get the layer output mutably.
sourcepub fn encoder_with_attention(&self) -> Option<&HiddenLayer>
pub fn encoder_with_attention(&self) -> Option<&HiddenLayer>
Get the output of an encoder layer.
Return a Some
value if the layer output is from an encoder layer,
or None
otherwise.