Cleaner test eval

Xilinx · Sep 24, 2024 · 346c155 · 346c155
1 parent bb32466
commit 346c155
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/tests/brevitas/graph/test_calibration.py b/tests/brevitas/graph/test_calibration.py
@@ -102,9 +102,8 @@ def forward(self, x):
 
     computed_scale = model.act.act_quant.scale(), model.act_1.act_quant.scale()
     reference_values = REFERENCE_SCALES[reference]
-    assert all([
-        torch.allclose(comp, torch.tensor(ref)) for comp,
-        ref in zip(computed_scale, reference_values)])
+    assert torch.allclose(computed_scale[0], torch.tensor(reference_values[0]))
+    assert torch.allclose(computed_scale[1], torch.tensor(reference_values[1]))
 
 
 def test_calibration_training_state():