Solved by caonv in post #4 
              
              
                Let’s assume that your CNN class looks like this: 
class CNN(nn.Module): 
    def __init__(self, *args, **kwargs):
        super(CNN, self).__init__()
        self.feature_extractor = ... # CNN layers or whatever
        self.classifier = .... # Linear layers
    
    def forward(self, x):
        x…