src/transformers/models/jukebox/modeling_jukebox.py/JukeboxPositionalEmbedding/__init__
class JukeboxPositionalEmbedding:    def __init__(self, embed_dim, width):
        super().__init__()
        self.pos_emb = nn.Parameter(torch.empty((embed_dim, width)))