marsggbo/wmt16_switch128_token_real_and_predicted_patterns

decode_pattern.shape: (num_layers, num_decoding_steps) predictor_pattern.shape: (num_decoding_steps, num_layers, top3_indices) We need to permute predictor_pattern via: predictor_pattern = predictor_pattern.permute(1,0,2) # (num_layers, num_decoding_steps, top3_indices)