Updates imagenet.ipynb to use GPU Colab runtime.
Unfortunately Colab TPU runtime only supports `jax==0.3.25`.
See also #2995, #2950, and #2861.
This PR updates the `examples/imagenet/imagenet.ipynb` Colab notebook to use the GPU runtime (and has some additional modifications to make it work with less host RAM).
Note: This PR depends on #3004.
