File "src/kd_train_text_to_image.py", line 790, in main
KeyError: 'up_blocks.0'

出问题的原因

dict
acts_tea读到dict
acts_stu没有读到dict
原因是
unet_teacher的结构后面直接接down_blocks（正常）
unet_teacher.down_blocks
但是unet的结构后面接的module之后才是down_blocks(不正常),
这个module事导致不能hook的原因
unet.module.down_blocks

解决办法

在student的索引里面添加module.