Solved by ptrblck in post #7 
              
              
                You could derive a custom class using the resnet class as its parent: 
import torchvision.models as models
from torchvision.models.resnet import ResNet, BasicBlock

class MyResNet18(ResNet):
    def __init__(self):
        super(MyResNet18, self).__init__(BasicBlock, [2, 2, 2, 2])
        
    def f…