Exctracting features from each frame of a video and stacking all the frames' features at the end

Question

I am working with videos. So, the input of my model has the shape of [batch_size,channels,frames,height,width]. So, my goal is to extract the the features from each frame separately using ResNet2D and also I expect to have a tensor with the same shape for my output. Thus, I wrote the following code and I am not sure if I designed it correctly when it comes to splitting the input data for feeding the ResNet2Dcand also I am not sure that using dim=2 is correct in torch.stack. Moreover, I think the code can become more efficient. Any idea is appreciated.

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


class my_model(nn.Module):
    def __init__(self, pretrained=True):
        super(my_model, self).__init__()
            
        self.featureExtractor = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)
        self.activation = {}
        def get_activation(name):
            def hook(model, input, output):
                self.activation[name] = output.detach()
            return hook

        self.featureExtractor.layer1.register_forward_hook(get_activation('layer1'))
        self.featureExtractor.layer2.register_forward_hook(get_activation('layer2'))
        self.featureExtractor.layer3.register_forward_hook(get_activation('layer3'))
        self.featureExtractor.layer4.register_forward_hook(get_activation('layer4'))
        
    def forward_features(self, x):   
        
        image_s = []
        features_1 = []
        features_2 = []
        features_3 = []
        features_4 = []

        for i in range(len_temporal):
            image_s = x[:,:,i:i+1,:,:]
            
            _ = self.featureExtractor(image_s.squeeze(2))
            layer1_output_Temporal_att = self.activation['layer1']
            features_1.append(layer1_output_Temporal_att)
            
            layer2_output_Temporal_att = self.activation['layer2']
            features_2.append(layer2_output_Temporal_att) 
            
            layer3_output_Temporal_att = self.activation['layer3']
            features_3.append(layer3_output_Temporal_att)   

            layer4_output_Temporal_att = self.activation['layer4']
            features_4.append(layer4_output_Temporal_att)                      
            
                       
            
        return torch.stack(features_1,dim=2), torch.stack(features_2,dim=2), torch.stack(features_3,dim=2), torch.stack(features_4,dim=2)




        
    def forward(self, x):
        
        Spatial_features_1, Spatial_features_2, Spatial_features_3, Spatial_features_4 = self.forward_features(x)   # torch.Size([16, 1024, 16, 14, 14])
                
        return Spatial_features_1, Spatial_features_2, Spatial_features_3, Spatial_features_4