Solved by Eta_C in post #2 
              
              
                import torch 
import torch.nn as nn 


class Extractor(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 4, 5)
        self.bn = nn.BatchNorm2d(4)
    
    def forward(self, x):
        return self.bn(self.conv)


class MyModel(nn.Module):
    def __init…