I am trying to replace a certain type of layer (e.g. nn.Conv2d) with my own custom variant in a way that I do not have to copy paste the model definition code and replace the relevant lines.
I tried to do this using the solution mentioned in: Replacing convs modules with custom convs, then NotImplementedError
However when using this with the resnet model from torchvision and trying to replace the nn.Conv2d layers it does not replace the conv2d layers in the downsample block, and when trying to use it with the alexnet model from torchvision it does not replace any conv2d layer.
What is the recommended way to replace a certain type of layer with another in such a way that the original code definition of the model does not need to be copy pasted and adapted?