Issue with unit tests on NVIdia A100 (GPU)

[DwarKapex]

Hi everyone.

I see the issue when run unit tests on NVidia A100 (GPU). Here is the link for more details.

Briefly:

=========================== short test summary info ============================
FAILED opt/gemma/gemma/layers_test.py::EinsumTest::test_rmsnorm0 - AssertionE...
FAILED opt/gemma/gemma/modules_test.py::FeedForwardTest::test_ffw0 - Assertio...
FAILED opt/gemma/gemma/positional_embeddings_test.py::PositionalEmbeddingsTest::test_adds_positional_embeddings0
FAILED opt/gemma/gemma/sampler_test.py::SamplerTest::test_forward_equivalence
================== 4 failed, 12 passed, 2 warnings in 26.55s ===================

The first 3 is similar to issues on V100 (#32), but the last one: 4. test_forward_equivalence link. Can you relax the tolerance when run on GPUs?