import torch
from transformer_decoder import CustomDecoderConfig, TransformerDecoder

config = CustomDecoderConfig()
model = TransformerDecoder(config)

input_ids = torch.randint(0, config.vocab_size, (2, 10))
encoder_output = torch.randn(2, 10, config.hidden_size)

logits = model(
    input_ids=input_ids,
    encoder_output=encoder_output
)

print("Input shape:", input_ids.shape)
print("Logits shape:", logits.shape)