Struct SpatialTransformer

pub struct SpatialTransformer {
    pub config: SpatialTransformerConfig,
    /* private fields */
}

Fields

config: SpatialTransformerConfig

Implementations

impl SpatialTransformer

pub fn new( vs: VarBuilder<'_>, in_channels: usize, n_heads: usize, d_head: usize, use_flash_attn: bool, config: SpatialTransformerConfig, ) -> Result<Self>

pub fn forward(&self, xs: &Tensor, context: Option<&Tensor>) -> Result<Tensor>

Trait Implementations

impl Debug for SpatialTransformer

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.