Notebook was run in Google Colab for testing it.
Indeed, when using CUDA, there should be a ".cpu()" call.
Took notice and modifying the notebook. Thank you!
Indeed, when using CUDA, there should be a ".cpu()" call.
Took notice and modifying the notebook. Thank you!