I think that the operation of obtaining a feature map is not simply a matter of dimensioning the network output. Because each feature map is transitioned through operations such as convolution.
def forward(self, x):
"""
feed forward
"""
y = self.densenet.features(x)
output = []
for stop in self.stdn_in:
output.append(y[:, :stop, :, :])