set keras backend to jax in classification_metrics_test (#57)

Author: jshin1394

src/metrax/classification_metrics_test.py

@@ -14,6 +14,9 @@
 
 """Tests for metrax classification metrics."""
 
+import os
+os.environ['KERAS_BACKEND'] = 'jax'
+
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax.numpy as jnp