Hi,
Thanks for your response. I need to clarify what I am trying to do.
In the sample code above, during each iteration of train(epoch), after updating the weights of the model as described, I want to use these weights or model.parameters() to “update” some of the features. i.e. features[:, f:] for some value f < features.shape[1]:
def train(epoch):
    # train model
    # update weights

    # update the last (features.shape[1] - f) features <--

for epoch in range(args.epochs):
    train(epoch)

The last (features.shape[1] - f) features need to be updated by computing the gradient of the error w.r.t. these features (given the recently computed model weights/parameters) and taking a gradient step.