Hi I’ve been trying to print out the different modules, layers and output sizes of each of them but I’ve been getting this problem:
** RuntimeError: running_mean should contain 3 elements not 64

after searching around I believe this is related to batchnorm but I’m not sure how to resolve it.
Minimal working example:
import torch
import torchvision.models as models

model = models.densenet201().cuda()
model.eval()
x = torch.FloatTensor(8, 3, 256, 256).cuda()
for feat in model.featuers:
     feat(x).size()

Any ideas?