Describe the bug
As the title says

To Reproduce
Use CUDA

<Paste your code here>

Expected behavior
Faster training with most of the GPU utilized

Additional context
na

Versions