Transforming pretrained ResNet50 to a hierachical model by dividing it into subnetworks

1

For extracting features from images I use pretrained ResNet50 (removing the classifier head). In order to transform the model to a hierarchical model, I want to use the following strategy:

import torch
from torchvision.models import resnet50, ResNet50_Weights
import torch.nn as nn

class model(nn.Module):
    def __init__(self):
        super(model, self).__init__()

        self.model = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)
        self.block1 = nn.Sequential(*list(self.model.children())[:5])
        self.block2 = nn.Sequential(*list(self.model.children())[5:6])
        self.block3 = nn.Sequential(*list(self.model.children())[6:7])
        self.block4 = nn.Sequential(*list(self.model.children())[7:8])


    def forward(self, x):   
        
        x1 = self.block1(x)    
        x2 = self.block2(x1)
        x3 = self.block3(x2)
        x4 = self.block4(x3)

        return(x1, x2, x3, x4)

this is because I need to have the intermediate outputs. My question is that does this approach cause any problem for assigning the pretrained weights to different layers (i.e., self.block2, self.block3 ,…)? Should I edit the source code for my goal or the mentioned strategy is a working idea?