Way to set the device type of tensor that comes out of pytorch distribution

Is there a way so set the tensor type that comes out of a distribution? For example, torch.OneHotCateogorical.to(“CUDA”) so that the tensor of categoricals that comes out is CUDA?