Hello,
I am trying to find a way to add layers to a model after without changing the definition of the model
For example let’s say i have this layer that I want to add after every ReLU layer i have in a ResNet
class SimpleCustomLayer(nn.Module):
    def __init__(self):
        pass

    def forward(self,x):
        log_tensor(x)
        return x



def add_module_after_layer(module:nn.Module, layer:nn.Module,model:nn.Module):
    for i in model.modules():
        if type(i) == type(layer):
            # Here i should add the module after the layer

    

model = ResNet()

add_module_after_layer(module= SimpleCustomLayer,
                       layer= torch.nn.ReLU,
                       model=model)

Any idea how to add this layer after layer i ?