Help to identify this layers in EfficientNet_b0

vision
1

I am implementing a paper that says:

we extract vectors from layers 7 (level 2), 20 (level 4), and 26 (level 5) in an EfficientNet-B5.

I am trying the torchvision.model.efficientnet_b0 implementation but, I am not sure what layers are.

Could you help me? I can not catch the exactly feature maps. Thanks

image

This is my best try

import torchvision

torch_model = torchvision.models.efficientnet_b0(False)
return_nodes = {
    'features.1': 'level2',
    'features.3': 'level4',
    'features.4': 'level5',
}
eff_model = create_feature_extractor(torch_model, return_nodes)
o = eff_model(torch.rand(1, 3, 224, 224))
for k in o:
  print(k, o[k].shape)

Output

level2 torch.Size([1, 16, 112, 112])
level4 torch.Size([1, 40, 28, 28])
level5 torch.Size([1, 80, 14, 14])