How to pass the user input as an object to a neural network?

Hi,
I’ve used hook function to extract the pertained weights at certain layer.

activation = {}
def get_activation(name):
    def hook(model, input, output):
        activation[name] = output.requires_grad_()
    return hook

Now I would like to make a user interface that the user will input a layer name (eg. “layer2”, “layer3”), and then pass it as an object to the model.
When the user input ‘layer2’, it will do
resnet_model.layer2.register_forward_hook(get_activation('layer2'))
When the user input ‘layer3’, it will do
resnet_model.layer3.register_forward_hook(get_activation('layer3'))

My problem is that I don’t know how to pass the string input by user as an object to a nn.

Any help and suggestions would be appreciated, thanks in advance.

You could add this argument to the __init__ method or a custom class method of your model:

class MyModel(nn.Module):
    def __init__(self, user_input=''):
        super(MyModel, self).__init__()
        print('user_input is {}'.format(user_input))
        # create the modules here and register the hook
    
    def my_custom_method(self, user_input=''):
        print('user_input is now {}'.format(user_input))
        
    def forward(self, x):
        return x

model = MyModel(user_input='a')
model.my_custom_method(user_input='b')