Don’t know about anything exactly like the PyTorch example, but you can probably put something together.
This thread could maybe be of interest
where the cuda docs are mentioned, where you can find a little about some basic data handling over gpus.