When you create an optimizer, you pass it the parameters you want it to optimize, so just give it the ones in the first layer.
Note also that when you call backwards, PyTorch will need to compute grads of variables other than the ones you specified when you created it, but they won’t be leaf variables.  When you call zero_grad, it will zero the grad of the parameters you specified when you created it.  It doesn’t need to zero the grad of the other variables it computed grad for because grads only accumulate in leaf variables.