about the tome with mlp module

[LWang1016]

Hi, @dbolya , a very solid work, cutrrently I apply tome with some diffusion-based models, and found that when perform on MLP module, even the pruning ratio is very small (0.1), the results become not correct, like this. Do you have some suggestions?

[dbolya]

It may be that the tomesd patch doesn't work on the model you're using. For instance, if that model implements the MLP differently or has extra modules. Basically the features got destroyed.

Some models may require custom implementations. You may want to check if they have implemented ToMe or a similar approach themselves.