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.