I have problem with decoder and task quite big to identify where is the error. Do you have any idea where do dig?
======================================================================
Decoder - Main Layers
======================================================================
Layer Type Parameters
----------------------------------------------------------------------
token_emb Embedding 1,280,000
pos_enc PositionalEncoding 0
dropout Dropout 0
transformer_decoder TransformerDecoder 2,372,352
output_projection Linear 1,285,000
----------------------------------------------------------------------
TOTAL 4,937,352
======================================================================
Expected Output
======================================================================
Decoder - Main Layers
======================================================================
Layer Type Parameters
----------------------------------------------------------------------
token_emb Embedding 1,280,000
pos_enc PositionalEncoding 0
dropout Dropout 0
transformer_decoder TransformerDecoder 2,372,352
output_projection Linear 1,285,000
----------------------------------------------------------------------
TOTAL 4,937,352
======================================================================
unittests.exercise_2(Decoder)
Failed test case: Failed forward pass test.
Expected: None
Got: shape '[1, 1, 8, 8]' is invalid for input of size 16