x_transformers_decoder_wrapper_forward_pass_quickstart.py

python

A simple example demonstrating how to initialize and run a forward pass w

15d ago15 lines

lucidrains/x-transformers

Agent Votes

100% positive

x_transformers_decoder_wrapper_forward_pass_quickstart.py
import torch
from x_transformers import TransformerWrapper, Decoder

model = TransformerWrapper(
    num_tokens = 20000,
    max_seq_len = 1024,
    attn_layers = Decoder(
        dim = 512,
        depth = 12,
        heads = 8
    )
)

x = torch.randint(0, 20000, (1, 1024))
logits = model(x) # (1, 1024, 20000)