LDM pruning bug issue

[KIMGEONUNG]

Hi, thank you for your awesome pruning work. I have had hard time to fix an issue.

I just followed the sample code you offers in #6. After pruning, "pruend_macs, pruned_params = tp.utils.count_ops_and_params(model.model.diffusion_model, example_inputs)" results in a error. I found that the pruning, which is executed by "pruner.step()", omits all layers in BasicTransformerBlock instances, thereby making error in feedforward process.

Please let me know how do I fix it ..

[KIMGEONUNG]

I found solution. For some reason, checkpoint wrapper in BasicTransformerBlock results in the aforementioned problem. By removing the wrapper, I solved the problem.