mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 04:46:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Update colab CPU test notebook to be more robust

Browse Source
This commit is contained in:
Jake VanderPlas 2021-07-13 09:30:56 -07:00
parent 208dd1ac3f
commit 5a0a46cbef
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tests/notebooks/colab_cpu.ipynb
View File

@ -95,7 +95,7 @@
"arr = jax.random.normal(key, (1000,))\n",
"device = arr.device_buffer.device()\n",
"print(f\"JAX device type: {device}\")\n",
"assert isinstance(device, xla_extension.CpuDevice), \"unexpected JAX device type\""
"assert device.platform == \"cpu\", f\"unexpected JAX device type: {device.platform}\""
],
"execution_count": 2,
"outputs": [