Hi @ladan_yang ,
Although there is no explicit error message here, however, the AssertionError is telling you that your code for sacled_do_product_attention() failed because the weights returned did not match those on the right, see line 62 from the output of the unit test.
Getting the correct dk is very important, I can see there is a need for your attention there.
Thank you so much! It is solved now!