Consider this. When you move your model to GPU, all rhe weights of the model are moved to GPU. Now when you have to do some computation on these weights that has to be done on the GPU because GPU-CPU-GPU transfers are expensive. So all the code of your model in turn runs on rhe GPU.