Ho,
Then PW = W.index_select(0, p), where, if W is a torch.cuda.XXXXTensor(), p has to be a torch.cuda.LongTensor.
If you already have some memory allocated for PW, you can use torch.index_select(W, 0, p, out=PW).