Hi @ptrblck,
I used your method as follows:
import os
import random

import torch
import torch.nn as nn
import torchvision
import torchvision.datasets as datasets
import torchvision.models as models
import torchvision.transforms as transforms

from utils import load_state_dict_from_url

######## To print layer outputs ########
class PrintLayer(nn.Module):
    def __init__(self):
        super(PrintLayer, self).__init__()
                    
    def forward(self, x):
        # Do your print / debug stuff here
        print(x)
        return x

########################################


######## AlexNet model ########

__all__ = ['AlexNet', 'alexnet']

model_urls = { 'alexnet': 'https://download.pytorch.org/models/alexnet-owt-4df8aa71.pth', }

class AlexNet(nn.Module):
    
    def __init__(self, num_classes=1000):
        super(AlexNet, self).__init__()
        self.features = nn.Sequential(
                nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
                PrintLayer(), # layer output
                nn.ReLU(inplace=True),
                nn.MaxPool2d(kernel_size=3, stride=2),
                nn.Conv2d(64, 192, kernel_size=5, padding=2),
                PrintLayer(), # layer output
                nn.ReLU(inplace=True),
                nn.MaxPool2d(kernel_size=3, stride=2),
                nn.Conv2d(192, 384, kernel_size=3, padding=1),
                PrintLayer(), # layer output
                nn.ReLU(inplace=True),
                nn.Conv2d(384, 256, kernel_size=3, padding=1),
                PrintLayer(), # layer output
                nn.ReLU(inplace=True),
                nn.Conv2d(256, 256, kernel_size=3, padding=1),
                PrintLayer(), # layer output
                nn.ReLU(inplace=True),
                nn.MaxPool2d(kernel_size=3, stride=2),
                )
        self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
        self.classifier = nn.Sequential(
                nn.Dropout(),
                nn.Linear(256 * 6 * 6, 4096),
                nn.ReLU(inplace=True),
                nn.Dropout(),
                nn.Linear(4096, 4096),
                nn.ReLU(inplace=True),
                nn.Linear(4096, num_classes),
                )
                                        
    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x


def alexnet(pretrained=False, progress=True, **kwargs):
    r"""AlexNet model architecture from the
    `"One weird trick..." <https://arxiv.org/abs/1404.5997>`_ paper.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    model = AlexNet(**kwargs)
    if pretrained:
        state_dict = load_state_dict_from_url(model_urls['alexnet'], progress=progress)
        model.load_state_dict(state_dict)
    return model

###############################


############# Load Data ##############

workers = 0
batchsize = 256

valdir = os.path.join( '/stor2/gakadam/GPU_RESEARCH/PyTorchFI/data/' + '/imagenet/', 'val')
normalize = transforms.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
val_loader = torch.utils.data.DataLoader(
        datasets.ImageFolder(
            valdir, transforms.Compose( [ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), normalize, ]),
            ),
        batch_size=batchsize, shuffle=False, num_workers=workers,)

model = alexnet(pretrained=True)

######################################


########### Classification ############

batch = next(iter(val_loader))
images, labels = batch

with torch.no_grad():
    output = model(images)
    out_soft = torch.nn.functional.softmax(output, dim=1)
    for i in out_soft:
        print(torch.argmax(i))

#######################################



But I get following error:
Traceback (most recent call last):
  File "inference_print_layer_output.py", line 104, in <module>
    model = alexnet(pretrained=True)
  File "inference_print_layer_output.py", line 85, in alexnet
    model.load_state_dict(state_dict)
  File "/home/USER/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1045, in load_state_dict
    self.__class__.__name__, "\n\t".join(error_msgs)))
RuntimeError: Error(s) in loading state_dict for AlexNet:
        Missing key(s) in state_dict: "features.4.weight", "features.4.bias", "features.11.weight", "features.11.bias", "features.14.weight", "features.14.bias".
        Unexpected key(s) in state_dict: "features.3.weight", "features.3.bias", "features.6.weight", "features.6.bias", "features.10.weight", "features.10.bias".
        size mismatch for features.8.weight: copying a param with shape torch.Size([256, 384, 3, 3]) from checkpoint, the shape in current model is torch.Size([384, 192, 3, 3]).
        size mismatch for features.8.bias: copying a param with shape torch.Size([256]) from checkpoint, the shape in current model is torch.Size([384]).

There is size mismatch even though PrintLayer does not change the tensor.
How to rectify this?