# Buggy code
def forward(
    self,
    src_tokens,
    src_lengths,
    cls_input: Optional[Tensor] = None,
    return_all_hiddens: bool = False,
):
    """
    Args:
        src_tokens (LongTensor): tokens in the source language of shape
            `(batch, src_len)`
        src_lengths (torch.LongTensor): lengths of each source sentence of
            shape `(batch)`
        return_all_hiddens (bool, optional): also return all of the
            intermediate hidden states (default: False).

    Returns:
        namedtuple:
            - **encoder_out** (Tensor): the last encoder layer's output of
              shape `(src_len, batch, embed_dim)`
            - **encoder_padding_mask** (ByteTensor): the positions of
              padding elements of shape `(batch, src_len)`
            - **encoder_embedding** (Tensor): the (scaled) embedding lookup
              of shape `(batch, src_len, embed_dim)`
            - **encoder_states** (List[Tensor]): all intermediate
              hidden states of shape `(src_len, batch, embed_dim)`.
              Only populated if *return_all_hiddens* is True.
    """
    if self.layer_wise_attention:
        return_all_hiddens = True

    x, encoder_embedding = self.forward_embedding(src_tokens)

    # B x T x C -> T x B x C
    x = x.transpose(0, 1)

    # compute padding mask
    encoder_padding_mask = src_tokens.eq(self.padding_idx)

    encoder_states = [] if return_all_hiddens else None

    # encoder layers
    for layer in self.layers:
        # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
        dropout_probability = torch.empty(1).uniform_()
        if not self.training or (dropout_probability > self.encoder_layerdrop):
            x = layer(x, encoder_padding_mask)
            if return_all_hiddens:
                assert encoder_states is not None
                encoder_states.append(x)

    if self.layer_norm is not None:
        x = self.layer_norm(x)
        if return_all_hiddens:
            encoder_states[-1] = x

    return EncoderOut(
        encoder_out=x,  # T x B x C
        encoder_padding_mask=encoder_padding_mask,  # B x T
        encoder_embedding=encoder_embedding,  # B x T x C
        encoder_states=encoder_states,  # List[T x B x C]
        src_tokens=None,
        src_lengths=None,
    )


# Traceback (IndexError)
bash train1.sh
2020-04-29 21:50:54 | INFO | fairseq.distributed_utils | distributed init (rank 1): tcp://localhost:14321
2020-04-29 21:50:54 | INFO | fairseq.distributed_utils | distributed init (rank 0): tcp://localhost:14321
2020-04-29 21:50:55 | INFO | fairseq.distributed_utils | initialized host zixi-MS-7B79 as rank 1
2020-04-29 21:50:55 | INFO | fairseq.distributed_utils | initialized host zixi-MS-7B79 as rank 0
2020-04-29 21:50:57 | INFO | fairseq_cli.train | Namespace(activation_dropout=0.1, activation_fn='relu', adam_betas='(0.9, 0.999)', adam_eps=1e-08, adaptive_input=False, adaptive_softmax_cutoff=None, adaptive_softmax_dropout=0, all_gather_list_size=16384, arch='transformer', attention_dropout=0.1, best_checkpoint_metric='ppl', bpe='subword_nmt', bpe_codes=None, bpe_separator='@@', broadcast_buffers=False, bucket_cap_mb=25, checkpoint_suffix='', clip_norm=25, cpu=False, criterion='cross_entropy', cross_self_attention=False, curriculum=0, data='temp/bpe/bin', data_buffer_size=0, dataset_impl=None, ddp_backend='no_c10d', decoder_attention_heads=4, decoder_embed_dim=256, decoder_embed_path=None, decoder_ffn_embed_dim=256, decoder_input_dim=256, decoder_layerdrop=0.1, decoder_layers=4, decoder_layers_to_keep=None, decoder_learned_pos=False, decoder_normalize_before=True, decoder_output_dim=256, device_id=0, disable_validation=False, distributed_backend='nccl', distributed_init_method='tcp://localhost:14321', distributed_no_spawn=False, distributed_port=-1, distributed_rank=0, distributed_world_size=2, distributed_wrapper='DDP', dropout=0.1, empty_cache_freq=0, encoder_attention_heads=4, encoder_embed_dim=256, encoder_embed_path=None, encoder_ffn_embed_dim=256, encoder_layerdrop=0.1, encoder_layers=4, encoder_layers_to_keep=None, encoder_learned_pos=False, encoder_normalize_before=True, eval_bleu=False, eval_bleu_args=None, eval_bleu_detok='space', eval_bleu_detok_args=None, eval_bleu_print_samples=False, eval_bleu_remove_bpe=None, eval_tokenized_bleu=False, fast_stat_sync=False, find_unused_parameters=False, fix_batches_to_gpus=False, fixed_validation_seed=None, force_anneal=None, fp16=False, fp16_init_scale=128, fp16_no_flatten_grads=False, fp16_scale_tolerance=0.0, fp16_scale_window=None, keep_best_checkpoints=10, keep_interval_updates=-1, keep_last_epochs=-1, layer_wise_attention=False, layernorm_embedding=False, left_pad_source='True', left_pad_target='False', load_alignments=False, localsgd_frequency=3, log_format=None, log_interval=100, lr=[0.005], lr_scheduler='fixed', lr_shrink=0.5, max_epoch=0, max_sentences=None, max_sentences_valid=None, max_source_positions=1024, max_target_positions=1024, max_tokens=5000, max_tokens_valid=5000, max_update=0, maximize_best_checkpoint_metric=False, memory_efficient_fp16=False, min_loss_scale=0.0001, min_lr=-1, model_parallel_size=1, no_cross_attention=False, no_epoch_checkpoints=False, no_last_checkpoints=False, no_progress_bar=False, no_save=False, no_save_optimizer_state=True, no_scale_embedding=False, no_token_positional_embeddings=False, nprocs_per_node=2, num_workers=1, optimizer='adam', optimizer_overrides='{}', patience=8, quant_noise_pq=0, quant_noise_pq_block_size=8, quant_noise_scalar=0, quantization_config_path=None, required_batch_size_multiple=8, reset_dataloader=False, reset_lr_scheduler=False, reset_meters=False, reset_optimizer=False, restore_file='checkpoint_last.pt', save_dir='checkpoints/transformer', save_interval=1, save_interval_updates=0, seed=1, sentence_avg=False, share_all_embeddings=False, share_decoder_input_output_embed=False, skip_invalid_size_inputs_valid_test=False, slowmo_algorithm='LocalSGD', slowmo_momentum=None, source_lang=None, target_lang=None, task='translation', tensorboard_logdir='', threshold_loss_scale=None, tokenizer=None, train_subset='train', truncate_source=False, update_freq=[8], upsample_primary=1, use_bmuf=False, use_old_adam=False, user_dir=None, valid_subset='valid', validate_interval=1, warmup_updates=0, weight_decay=0.0)
2020-04-29 21:50:57 | INFO | fairseq.tasks.translation | [src] dictionary: 48947 types
2020-04-29 21:50:57 | INFO | fairseq.tasks.translation | [tgt] dictionary: 48613 types
2020-04-29 21:50:57 | INFO | fairseq.data.data_utils | loaded 8204 examples from: temp/bpe/bin/valid.src-tgt.src
2020-04-29 21:50:57 | INFO | fairseq.data.data_utils | loaded 8204 examples from: temp/bpe/bin/valid.src-tgt.tgt
2020-04-29 21:50:57 | INFO | fairseq.tasks.translation | temp/bpe/bin valid src-tgt 8204 examples
2020-04-29 21:50:58 | INFO | fairseq_cli.train | TransformerModel(
(encoder): TransformerEncoder(
(embed_tokens): Embedding(48947, 256, padding_idx=1)
(embed_positions): SinusoidalPositionalEmbedding()
(layers): ModuleList(
(0): TransformerEncoderLayer(
(self_attn): MultiheadAttention(
(k_proj): Linear(in_features=256, out_features=256, bias=True)
(v_proj): Linear(in_features=256, out_features=256, bias=True)
(q_proj): Linear(in_features=256, out_features=256, bias=True)
(out_proj): Linear(in_features=256, out_features=256, bias=True)
)
(self_attn_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
(fc1): Linear(in_features=256, out_features=256, bias=True)
(fc2): Linear(in_features=256, out_features=256, bias=True)
(final_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
)
(1): TransformerEncoderLayer(
(self_attn): MultiheadAttention(
(k_proj): Linear(in_features=256, out_features=256, bias=True)
(v_proj): Linear(in_features=256, out_features=256, bias=True)
(q_proj): Linear(in_features=256, out_features=256, bias=True)
(out_proj): Linear(in_features=256, out_features=256, bias=True)
)
(self_attn_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
(fc1): Linear(in_features=256, out_features=256, bias=True)
(fc2): Linear(in_features=256, out_features=256, bias=True)
(final_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
)
(2): TransformerEncoderLayer(
(self_attn): MultiheadAttention(
(k_proj): Linear(in_features=256, out_features=256, bias=True)
(v_proj): Linear(in_features=256, out_features=256, bias=True)
(q_proj): Linear(in_features=256, out_features=256, bias=True)
(out_proj): Linear(in_features=256, out_features=256, bias=True)
)
(self_attn_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
(fc1): Linear(in_features=256, out_features=256, bias=True)
(fc2): Linear(in_features=256, out_features=256, bias=True)
(final_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
)
(3): TransformerEncoderLayer(
(self_attn): MultiheadAttention(
(k_proj): Linear(in_features=256, out_features=256, bias=True)
(v_proj): Linear(in_features=256, out_features=256, bias=True)
(q_proj): Linear(in_features=256, out_features=256, bias=True)
(out_proj): Linear(in_features=256, out_features=256, bias=True)
)
(self_attn_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
(fc1): Linear(in_features=256, out_features=256, bias=True)
(fc2): Linear(in_features=256, out_features=256, bias=True)
(final_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
)
)
(layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
)
(decoder): TransformerDecoder(
(embed_tokens): Embedding(48613, 256, padding_idx=1)
(embed_positions): SinusoidalPositionalEmbedding()
(layers): ModuleList(
(0): TransformerDecoderLayer(
(self_attn): MultiheadAttention(
(k_proj): Linear(in_features=256, out_features=256, bias=True)
(v_proj): Linear(in_features=256, out_features=256, bias=True)
(q_proj): Linear(in_features=256, out_features=256, bias=True)
(out_proj): Linear(in_features=256, out_features=256, bias=True)
)
(self_attn_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
(encoder_attn): MultiheadAttention(
(k_proj): Linear(in_features=256, out_features=256, bias=True)
(v_proj): Linear(in_features=256, out_features=256, bias=True)
(q_proj): Linear(in_features=256, out_features=256, bias=True)
(out_proj): Linear(in_features=256, out_features=256, bias=True)
)
(encoder_attn_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
(fc1): Linear(in_features=256, out_features=256, bias=True)
(fc2): Linear(in_features=256, out_features=256, bias=True)
(final_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
)
(1): TransformerDecoderLayer(
(self_attn): MultiheadAttention(
(k_proj): Linear(in_features=256, out_features=256, bias=True)
(v_proj): Linear(in_features=256, out_features=256, bias=True)
(q_proj): Linear(in_features=256, out_features=256, bias=True)
(out_proj): Linear(in_features=256, out_features=256, bias=True)
)
(self_attn_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
(encoder_attn): MultiheadAttention(
(k_proj): Linear(in_features=256, out_features=256, bias=True)
(v_proj): Linear(in_features=256, out_features=256, bias=True)
(q_proj): Linear(in_features=256, out_features=256, bias=True)
(out_proj): Linear(in_features=256, out_features=256, bias=True)
)
(encoder_attn_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
(fc1): Linear(in_features=256, out_features=256, bias=True)
(fc2): Linear(in_features=256, out_features=256, bias=True)
(final_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
)
(2): TransformerDecoderLayer(
(self_attn): MultiheadAttention(
(k_proj): Linear(in_features=256, out_features=256, bias=True)
(v_proj): Linear(in_features=256, out_features=256, bias=True)
(q_proj): Linear(in_features=256, out_features=256, bias=True)
(out_proj): Linear(in_features=256, out_features=256, bias=True)
)
(self_attn_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
(encoder_attn): MultiheadAttention(
(k_proj): Linear(in_features=256, out_features=256, bias=True)
(v_proj): Linear(in_features=256, out_features=256, bias=True)
(q_proj): Linear(in_features=256, out_features=256, bias=True)
(out_proj): Linear(in_features=256, out_features=256, bias=True)
)
(encoder_attn_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
(fc1): Linear(in_features=256, out_features=256, bias=True)
(fc2): Linear(in_features=256, out_features=256, bias=True)
(final_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
)
(3): TransformerDecoderLayer(
(self_attn): MultiheadAttention(
(k_proj): Linear(in_features=256, out_features=256, bias=True)
(v_proj): Linear(in_features=256, out_features=256, bias=True)
(q_proj): Linear(in_features=256, out_features=256, bias=True)
(out_proj): Linear(in_features=256, out_features=256, bias=True)
)
(self_attn_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
(encoder_attn): MultiheadAttention(
(k_proj): Linear(in_features=256, out_features=256, bias=True)
(v_proj): Linear(in_features=256, out_features=256, bias=True)
(q_proj): Linear(in_features=256, out_features=256, bias=True)
(out_proj): Linear(in_features=256, out_features=256, bias=True)
)
(encoder_attn_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
(fc1): Linear(in_features=256, out_features=256, bias=True)
(fc2): Linear(in_features=256, out_features=256, bias=True)
(final_layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
)
)
(layer_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)
(output_projection): Linear(in_features=256, out_features=48613, bias=False)
)
)
2020-04-29 21:50:58 | INFO | fairseq_cli.train | model transformer, criterion CrossEntropyCriterion
2020-04-29 21:50:58 | INFO | fairseq_cli.train | num. model params: 41642240 (num. trained: 41642240)
2020-04-29 21:50:58 | INFO | fairseq_cli.train | training on 2 GPUs
2020-04-29 21:50:58 | INFO | fairseq_cli.train | max tokens per GPU = 5000 and max sentences per GPU = None
2020-04-29 21:50:58 | INFO | fairseq.trainer | no existing checkpoint found checkpoints/transformer/checkpoint_last.pt
2020-04-29 21:50:58 | INFO | fairseq.trainer | loading train data for epoch 1
2020-04-29 21:50:58 | INFO | fairseq.data.data_utils | loaded 1130841 examples from: temp/bpe/bin/train.src-tgt.src
2020-04-29 21:50:58 | INFO | fairseq.data.data_utils | loaded 1130841 examples from: temp/bpe/bin/train.src-tgt.tgt
2020-04-29 21:50:58 | INFO | fairseq.tasks.translation | temp/bpe/bin train src-tgt 1130841 examples
2020-04-29 21:51:05 | INFO | fairseq.trainer | NOTE: your device may support faster training with --fp16
epoch 001:   0%|                                                                                                                                                    | 0/210 [00:00<?, ?it/s]/opt/conda/conda-bld/pytorch_1587428398394/work/torch/csrc/utils/python_arg_parser.cpp:756: UserWarning: This overload of add_ is deprecated:
add_(Number alpha, Tensor other)
Consider using one of the following signatures instead:
add_(Tensor other, *, Number alpha)
/opt/conda/conda-bld/pytorch_1587428398394/work/torch/csrc/utils/python_arg_parser.cpp:756: UserWarning: This overload of add_ is deprecated:
add_(Number alpha, Tensor other)
Consider using one of the following signatures instead:
add_(Tensor other, *, Number alpha)
epoch 001 | valid on 'valid' subset | loss 4.067 | ppl 16.77 | wps 119770 | wpb 6793.9 | bsz 341.8 | num_updates 210
epoch 001 | valid on 'valid' subset | loss 4.067 | ppl 16.77 | wps 119685 | wpb 6793.9 | bsz 341.8 | num_updates 210
epoch 001 | loss 6.743 | ppl 107.11 | wps 59759.2 | ups 0.79 | wpb 76004.5 | bsz 5385 | num_updates 210 | lr 0.005 | gnorm 0.564 | clip 0 | train_wall 236 | wall 276
epoch 002:   0%|                                                                                                                                                    | 0/210 [00:00<?, ?it/s]2020-04-29 21:55:35 | INFO | fairseq.checkpoint_utils | saved checkpoint checkpoints/transformer/checkpoint1.pt (epoch 1 @ 210 updates, score 16.77) (writing took 0.49035206900043704 seconds)
epoch 001 | loss 6.743 | ppl 107.11 | wps 59649.1 | ups 0.78 | wpb 76004.5 | bsz 5385 | num_updates 210 | lr 0.005 | gnorm 0.564 | clip 0 | train_wall 236 | wall 276
epoch 002 | valid on 'valid' subset | loss 1.464 | ppl 2.76 | wps 121245 | wpb 6793.9 | bsz 341.8 | num_updates 420 | best_ppl 2.76
epoch 002 | valid on 'valid' subset | loss 1.464 | ppl 2.76 | wps 118177 | wpb 6793.9 | bsz 341.8 | num_updates 420 | best_ppl 2.76
epoch 002 | loss 1.986 | ppl 3.96 | wps 57927.6 | ups 0.76 | wpb 76004.5 | bsz 5385 | num_updates 420 | lr 0.005 | gnorm 0.262 | clip 0 | train_wall 243 | wall 551
epoch 003:   0%|                                                                                                                                                    | 0/210 [00:00<?, ?it/s]2020-04-29 22:00:13 | INFO | fairseq.checkpoint_utils | saved checkpoint checkpoints/transformer/checkpoint2.pt (epoch 2 @ 420 updates, score 2.76) (writing took 3.3108816809999553 seconds)
epoch 002 | loss 1.986 | ppl 3.96 | wps 57341.3 | ups 0.75 | wpb 76004.5 | bsz 5385 | num_updates 420 | lr 0.005 | gnorm 0.262 | clip 0 | train_wall 243 | wall 555
epoch 003 | valid on 'valid' subset | loss 1.184 | ppl 2.27 | wps 120108 | wpb 6793.9 | bsz 341.8 | num_updates 630 | best_ppl 2.27
epoch 003 | valid on 'valid' subset | loss 1.184 | ppl 2.27 | wps 117797 | wpb 6793.9 | bsz 341.8 | num_updates 630 | best_ppl 2.27
epoch 003 | loss 1.224 | ppl 2.34 | wps 57350.6 | ups 0.75 | wpb 76004.5 | bsz 5385 | num_updates 630 | lr 0.005 | gnorm 0.137 | clip 0 | train_wall 246 | wall 830
epoch 004:   0%|                                                                                                                                                    | 0/210 [00:00<?, ?it/s]2020-04-29 22:04:51 | INFO | fairseq.checkpoint_utils | saved checkpoint checkpoints/transformer/checkpoint3.pt (epoch 3 @ 630 updates, score 2.27) (writing took 3.109906989000592 seconds)
epoch 003 | loss 1.224 | ppl 2.34 | wps 57392.5 | ups 0.76 | wpb 76004.5 | bsz 5385 | num_updates 630 | lr 0.005 | gnorm 0.137 | clip 0 | train_wall 243 | wall 833
Traceback (most recent call last):
File "train.py", line 11, in <module>
cli_main()
File "/home/zixi/EE-599/fairseq/fairseq_cli/train.py", line 355, in cli_main
nprocs=args.distributed_world_size,
File "/home/zixi/anaconda3/envs/ROC/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 200, in spawn
return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
File "/home/zixi/anaconda3/envs/ROC/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 158, in start_processes
while not context.join():
File "/home/zixi/anaconda3/envs/ROC/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 119, in join
raise Exception(msg)
Exception:

-- Process 1 terminated with the following error:
Traceback (most recent call last):
File "/home/zixi/anaconda3/envs/ROC/lib/python3.7/site-packages/torch/multiprocessing/spawn.py", line 20, in _wrap
fn(i, *args)
File "/home/zixi/EE-599/fairseq/fairseq_cli/train.py", line 324, in distributed_main
main(args, init_distributed=True)
File "/home/zixi/EE-599/fairseq/fairseq_cli/train.py", line 117, in main
valid_losses = train(args, trainer, task, epoch_itr, max_update)
File "/home/zixi/anaconda3/envs/ROC/lib/python3.7/contextlib.py", line 74, in inner
return func(*args, **kwds)
File "/home/zixi/EE-599/fairseq/fairseq_cli/train.py", line 187, in train
log_output = trainer.train_step(samples)
File "/home/zixi/anaconda3/envs/ROC/lib/python3.7/contextlib.py", line 74, in inner
return func(*args, **kwds)
File "/home/zixi/EE-599/fairseq/fairseq/trainer.py", line 379, in train_step
ignore_grad=is_dummy_batch,
File "/home/zixi/EE-599/fairseq/fairseq/tasks/fairseq_task.py", line 341, in train_step
loss, sample_size, logging_output = criterion(model, sample)
File "/home/zixi/anaconda3/envs/ROC/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in __call__
result = self.forward(*input, **kwargs)
File "/home/zixi/EE-599/fairseq/fairseq/criterions/cross_entropy.py", line 29, in forward
net_output = model(**sample['net_input'])
File "/home/zixi/anaconda3/envs/ROC/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in __call__
result = self.forward(*input, **kwargs)
File "/home/zixi/EE-599/fairseq/fairseq/legacy_distributed_data_parallel.py", line 86, in forward
return self.module(*inputs, **kwargs)
File "/home/zixi/anaconda3/envs/ROC/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in __call__
result = self.forward(*input, **kwargs)
File "/home/zixi/EE-599/fairseq/fairseq/models/transformer.py", line 272, in forward
return_all_hiddens=return_all_hiddens,
File "/home/zixi/anaconda3/envs/ROC/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in __call__
result = self.forward(*input, **kwargs)
File "/home/zixi/EE-599/fairseq/fairseq/models/transformer.py", line 498, in forward
encoder_states[-1] = x
IndexError: list assignment index out of range

# Fixed code
def forward(
    self,
    src_tokens,
    src_lengths,
    cls_input: Optional[Tensor] = None,
    return_all_hiddens: bool = False,
):
    """
    Args:
        src_tokens (LongTensor): tokens in the source language of shape
            `(batch, src_len)`
        src_lengths (torch.LongTensor): lengths of each source sentence of
            shape `(batch)`
        return_all_hiddens (bool, optional): also return all of the
            intermediate hidden states (default: False).

    Returns:
        namedtuple:
            - **encoder_out** (Tensor): the last encoder layer's output of
              shape `(src_len, batch, embed_dim)`
            - **encoder_padding_mask** (ByteTensor): the positions of
              padding elements of shape `(batch, src_len)`
            - **encoder_embedding** (Tensor): the (scaled) embedding lookup
              of shape `(batch, src_len, embed_dim)`
            - **encoder_states** (List[Tensor]): all intermediate
              hidden states of shape `(src_len, batch, embed_dim)`.
              Only populated if *return_all_hiddens* is True.
    """
    x, encoder_embedding = self.forward_embedding(src_tokens)

    # B x T x C -> T x B x C
    x = x.transpose(0, 1)

    # compute padding mask
    encoder_padding_mask = src_tokens.eq(self.padding_idx)

    encoder_states = [] if return_all_hiddens else None

    # encoder layers
    for layer in self.layers:
        x = layer(x, encoder_padding_mask)
        if return_all_hiddens:
            assert encoder_states is not None
            encoder_states.append(x)

    if self.layer_norm is not None:
        x = self.layer_norm(x)

    return EncoderOut(
        encoder_out=x,  # T x B x C
        encoder_padding_mask=encoder_padding_mask,  # B x T
        encoder_embedding=encoder_embedding,  # B x T x C
        encoder_states=encoder_states,  # List[T x B x C]
        src_tokens=None,
        src_lengths=None,
    )