W2 UNQ_C8 Error in Training loop

[NV1]

Thank you @yhuang.
The following thread helped with determining the parameters @yhuang was referring to:

C4W2 UNQ_C8 training_loop error

Nevermind, I have solved this problem. It originates from UNQ_C5, where I mistaken set the d_head as d_feature for both the compute_attention_heads_closure and compute_attention_output_closure functions. It has been resolved.