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

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
When the user input ‘layer3’, it will do

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')