reproduce DiffCR results

[LWang1016]

Hi, @ranery , This work is very promising ,and currently I apply DiffCR to Sana, which is a linear attention-based model(https://github.com/NVlabs/Sana), and now I only test the routing with a fix compression ration for all the layer, , and after the training, the results is noisy. Can you give me some suggestion?
Here is my code (same with your pesudo code.)

[ranery]

How long have you been fine-tuning? One way to debug is to visualize the router’s prediction maps to check whether it has learned useful patterns.