Thanks for the code.
You are creating some layers inside the forward method on the fly.
These layers will be reinitialized in each forward pass, so they won’t be trained, and also won’t get pushed to the GPU, since they are unknown to the model when you call model.to(device).
The standard way would be to create these layers in the __init__ method (as your other layers) and just apply them on the activation in forward.
If you really want to use randomly initialized layers in the forward pass, you could use your code and push these newly created layers manually using the device attribute of another already registered parameter.