Only can run in CPU
miss-rain opened this issue · 2 comments
miss-rain commented
Both in Google Colab with TPU and my own Unbuntu with Cuda,
your code only run in CPU, not run in GPU or TPU!
zizhaozhang commented
@miss-rain I think it is related to your GPU type and cuda version and whether are them compatible with the installed jax version. It is not related to this code particularly.
Could you pls first make sure you can install jax and use jax to check GPU is detected (using this google/jax#971).
KingSpencer commented
Please refer to issue #1.