EmbeddingBlock

class ssm.model.block.embedding_block.EmbeddingBlock(model, vocab_size, model_dim, mem_tokens, out_dim=None)

Bases: Module

forward(x)

Forward pass of the embedding block.

Parameters:

x (torch.Tensor) – The input tensor.

Returns:

The output tensor.

Return type:

torch.Tensor