Based on your description, it seems you would like to create a weight parameter, which will scale the inputs, and which should be trainable.
If that’s the case, you could initialize it as:
weight = nn.Parameter(torch.randn(3).to(device))

and pass it with all other parameters to the optimizer:
optimizer = optim.Adam([weight] + list(model.parameters()), lr=1e-3)

This will make sure to optimizer all passed parameters.
Autograd will take care of creating the computation graph based on the applied operations.