You could either use forward hooks to get the desired activation as described here or replace the unwanted layers with nn.Identity modules.
You could either use forward hooks to get the desired activation as described here or replace the unwanted layers with nn.Identity modules.