Fix keras tests for GPU based docker images

[TRNWWZ]

Issue: keras tests are failing when running it for GPU .

More background: We hard code keras to a legacy commit: https://github.com/aws/sagemaker-distribution/blob/main/test/test_artifacts/v1/keras.test.Dockerfile#L19 due to TensorFlow version conflict issue, but with upgrading to new minor version.

After removing this line and rerun the test, test failed again complaining missing dependency like jax.

Acceptance Criteria: keras unit tests should pass.

Image Tags

1.8.0-gpu