diff --git a/.github/workflows/UnitTests.yml b/.github/workflows/UnitTests.yml
index f858a2532..6024553a2 100644
--- a/.github/workflows/UnitTests.yml
+++ b/.github/workflows/UnitTests.yml
@@ -56,3 +56,7 @@ jobs:
       run: |
         source venv/bin/activate
         python3 MaxText/decode.py MaxText/configs/base.yml run_name=runner_$(date +%Y-%m-%d-%H-%M) base_output_directory=gs://runner-maxtext-logs dataset_path=gs://maxtext-dataset steps=2
+    - name: Test int8_training
+      run: |
+        source venv/bin/activate
+        python3 MaxText/train.py MaxText/configs/base.yml run_name=runner_$(date +%Y-%m-%d-%H-%M) base_output_directory=gs://runner-maxtext-logs dataset_path=gs://maxtext-dataset int8_training=true steps=2
diff --git a/MaxText/aqt/jax/v2/aqt_dot_general.py b/MaxText/aqt/jax/v2/aqt_dot_general.py
index 5676a0823..f17a56a2c 100644
--- a/MaxText/aqt/jax/v2/aqt_dot_general.py
+++ b/MaxText/aqt/jax/v2/aqt_dot_general.py
@@ -540,10 +540,10 @@ def ret_dg(
     assert (
         precision is None
     ), f'Precision {precision} requested together with quantization.'
-    assert preferred_element_type is None, (
-        f'Preferred_element_typerecision {preferred_element_type} requested'
-        ' together with quantization.'
-    )
+
+    # The quantized einsum ignores preferred_element_type (b/302728979)
+    preferred_element_type = None
+
     assert lhs.dtype == rhs.dtype, (
         'The only reason we need that, is because we need to determine return'
         ' type.'