Relax tolerance in attention_projection_fq_test

PiperOrigin-RevId: 607472600
google · Feb 15, 2024 · 54d6d86 · 54d6d86
1 parent 5605af8
commit 54d6d86
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/praxis/layers/quantization/tests/attention_projection_fq_test.py b/praxis/layers/quantization/tests/attention_projection_fq_test.py
@@ -351,7 +351,7 @@ def loss(params, inputs, targets):
     self.assertAllClose(updated_w_f_tensor, updated_w_q_tensor, atol=1e-3)
 
     # 2. Value check.
-    self.assertNestedListClose(updated_w_q, expected_trained_weight)
+    self.assertNestedListClose(updated_w_q, expected_trained_weight, places=2)
 
   # Test the training with FQ quantization.
   @parameterized.parameters(