How to record the mapping relationship between remaining channels and the origin channels?

[MqLeet]

Hi @VainF , thanks to your excellent work! I am trying to prune DDPM and I want to record the mapping relationship between remaining channels and the origin channels.

For example: a Conv2d from (Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1)) =》 (Conv2d(224, 96, kernel_size=(1, 1), stride=(1, 1)), the idx=255 in output channel may become the idx=223 after pruning, how can I record the mapping relationship?

Looking forward to you reply! Best wishes, Qianli