[BUG] How to load checkpoint in NRMS/NAML model?

[kiennt120]

I try to load checkpoint by using argument callbacks but I get that problem

Can you show me another way to load the checkpoint work fine. Thank you!