jax and whisper-jax support running on GPUs in addition to TPUs, this particular demo just runs on a TPU.