def forward(self, x):
        for layers in self.front_process:
            print(layers)
            if isinstance(layers, nn.MaxPool2d):
                print('\ngot target1\n')
            
        print('\n\nmiddle \n\n')
                      
        for layers in self.back_process:
            print(layers)
            if isinstance(layers, nn.MaxUnpool2d):
                print('\ngot target2\n')

You can try this to extract maxpool and maxunpool, but I don’t know whether it will influence the speed or not, since I’m quite new to pytorch.