Enable AMDGPU tests in CI #206
Description
Thanks to #205 we now have GPU tests as part of CI, but currently we have crashes on the AMDGPU runner. It is quite possible the crashes are simply due to incorrectly setting up the runner.
It would be very valuable if someone with an AMD gpu runs the tests and confirms they work. Next step would be to create a docker file in which the tests still succeed. Then that docker file can be used to make a buildkite runner.
The current runner is docker.io/rocm/jax. The rocm-terminal container is not sufficient as it does not have all rocm libraries.