Views
No views yet
1import torch
2from transformer_decoder import CustomDecoderConfig, TransformerDecoder
3
4config = CustomDecoderConfig()
5model = TransformerDecoder(config)
6
7input_ids = torch.randint(0, config.vocab_size, (2, 10))
8encoder_output = torch.randn(2, 10, config.hidden_size)
9
10logits = model(input_ids=input_ids, encoder_output=encoder_output)
11print(logits.shape)torch.Size([2, 10, 30522])