Given the transformer architecture below, what is the embed_dim as described in other literatures?

import torch
import torch.nn as nn
import torch.nn.functional as F
import math


class RMSNorm(nn.Module):
    def __init__(self, d_model, eps=1e-6):
        super(RMSNorm, self).__init__()
        self.eps = eps
        self.scale = nn.Parameter(torch.ones(d_model))

    def forward(self, x):
        norm = x.norm(2, dim=-1, keepdim=True)
        return self.scale * x / (norm + self.eps)


class SwiGLU(nn.Module):
    def __init__(self, d_model, d_ff):
        super(SwiGLU, self).__init__()
        self.linear1 = nn.Linear(d_model, d_ff * 2)
        self.linear2 = nn.Linear(d_ff, d_model)

    def forward(self, x):
        x_proj = self.linear1(x)
        x1, x2 = x_proj.chunk(2, dim=-1)
        x = x1 * F.silu(x2)
        return self.linear2(x)


class ScaledDotProductAttention(nn.Module):
    def __init__(self, d_k, dropout=0.1):
        super(ScaledDotProductAttention, self).__init__()
        self.sqrt_dk = math.sqrt(d_k)
        self.dropout = nn.Dropout(dropout)

    def forward(self, q, k, v, mask=None):
        scores = torch.matmul(q, k.transpose(-2, -1)) / self.sqrt_dk
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        attn = self.dropout(F.softmax(scores, dim=-1))
        output = torch.matmul(attn, v)
        return output


class MultiheadAttention(nn.Module):
    def __init__(self, d_model, n_heads, dropout=0.1):
        super(MultiheadAttention, self).__init__()
        self.d_k = d_model // n_heads
        self.n_heads = n_heads

        self.q_linear = nn.Linear(d_model, d_model)
        self.k_linear = nn.Linear(d_model, d_model)
        self.v_linear = nn.Linear(d_model, d_model)
        self.out_linear = nn.Linear(d_model, d_model)
        self.attention = ScaledDotProductAttention(self.d_k, dropout)

    def forward(self, q, k, v, mask=None):
        batch_size = q.size(0)

        q = self.q_linear(q).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        k = self.k_linear(k).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
        v = self.v_linear(v).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)

        if mask is not None:
            mask = mask.unsqueeze(1)

        attn_output = self.attention(q, k, v, mask)

        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.n_heads * self.d_k)
        output = self.out_linear(attn_output)

        return output


class TransformerBlock(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super(TransformerBlock, self).__init__()
        self.n_heads = n_heads

        self.norm1 = RMSNorm(d_model)
        self.norm2 = RMSNorm(d_model)

        self.attn = MultiheadAttention(d_model, n_heads, dropout)
        self.ff = SwiGLU(d_model, d_ff)

        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, x, mask=True):
        x2 = self.norm1(x)
        attn_output = self.attn(x2, x2, x2, mask)
        x = x + self.dropout1(attn_output)
        
        x2 = self.norm2(x)
        ff_output = self.ff(x2)
        x = x + self.dropout2(ff_output)
        return x


class Transformer(nn.Module):
    def __iinput_size = (1, n_channels, d_model) nit__(self, d_model, n_heads, n_layers, d_ff, n_channels, max_len, dropout=0.1):
        super(Transformer, self).__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.n_layers = n_layers

        self.embedding = nn.Linear(n_channels, d_model)
        self.pos_embedding = nn.Parameter(torch.zeros(1, max_len, d_model))

        self.layers = nn.ModuleList([
            TransformerBlock(d_model, n_heads, d_ff, dropout) for _ in range(n_layers)
        ])

        self.fc_out = nn.Linear(d_model, n_channels)

    def forward(self, x, mask=True):
        batch_size, n_channels, seq_len = x.size()
        x = self.embedding(x.transpose(1, 2)) + self.pos_embedding[:, :seq_len, :]  # Project and add positional embeddings

        if mask:
            mask = self.generate_mask(x)

        for layer in self.layers:
            x = layer(x, mask)

        x = self.fc_out(x)
        return x.transpose(1, 2)  # Transpose back to original format (batch_size, n_channels, seq_len)

    def generate_mask(self, x):
        seq_len = x.size(-1)
        mask = torch.tril(torch.ones(seq_len, seq_len)).to(x.device)
        return mask.unsqueeze(0).unsqueeze(1)  # For compatibility with multihead attention


d_model = 256
n_heads = 4
d_ff = 1024
n_layers = 2
n_channels = 32
max_len = 256
dropout = 0.1

model = Transformer(
    d_model, 
    n_heads, 
    n_layers, 
    d_ff, 
    n_channels, 
    max_len, 
    dropout
)


from torchinfo import summary
input_size = (1, n_channels, d_model) 
print(summary(model, input_size))