Reduce required accuracy in ranking_agent_test.

PiperOrigin-RevId: 703085807 Change-Id: Iefd35170f48a9a8c6ce8f5b44d5e2dc32f6b3f24
tensorflow · Dec 5, 2024 · ee4d6fe · ee4d6fe
1 parent a809aa8
commit ee4d6fe
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tf_agents/bandits/agents/ranking_agent_test.py b/tf_agents/bandits/agents/ranking_agent_test.py
@@ -463,7 +463,9 @@ def testPositionalBiasParams(
     agent.train(experience)
     weights = agent._construct_sample_weights(scores, observations, None)
     self.assertAllEqual(weights.shape, [batch_size, num_slots])
-    self.assertAllClose(weights[-1, 1], expected_second_weight)
+    self.assertAllClose(
+        weights[-1, 1], expected_second_weight, atol=1e-3, rtol=1e-3
+    )
 
 
 if __name__ == '__main__':