SmallNet-EM-routing is slower than Resnet20-EM-routing

[humbleman24]

Hi, thanks for your great work!

I’m currently trying to replicate the EM Routing experiments. I noticed that in the original paper, the authors used shallower convolutional layers, so I switched from ResNet20 (which you implemented) to SmallNet.

However, I found something strange:
On RTX 4090, training with ResNet20 takes only about 30 seconds per epoch, but training with SmallNet takes about 2 minutes per epoch—which is much slower, even though SmallNet is supposed to have less computation.

Do you have any idea why this might be happening? Is there any optimization I should consider when using SmallNet, or is there something specific in the implementation that could cause this slowdown?

Thanks for your help!