fix tests

Xilinx · Sep 12, 2024 · 16d75d8 · 16d75d8
1 parent 6e0c4fc
commit 16d75d8
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/src/brevitas/quant_tensor/int_quant_tensor.py b/src/brevitas/quant_tensor/int_quant_tensor.py
@@ -15,7 +15,7 @@
 from .torch_handler import QUANT_TENSOR_FN_HANDLER
 
 IS_VALID_ATOL = 2e-1
-BFLOAT16_IS_VALID_ATOL = 0.5
+B_FLOAT16_IS_VALID_ATOL = 0.5
 
 
 class IntQuantTensor(IntQuantTensorBase, QuantTensor):
@@ -78,7 +78,8 @@ def is_valid(self):
             pre_round_int_value = self._pre_round_int_value
             rounded_int_value = torch.round(pre_round_int_value)
             max_abs_diff = torch.max(torch.abs(pre_round_int_value - rounded_int_value))
-            atol = BFLOAT16_IS_VALID_ATOL if self.value.dtype == torch.bfloat16 else IS_VALID_ATOL
+            atol = B_FLOAT16_IS_VALID_ATOL if self.value.dtype in (
+                torch.bfloat16, torch.float16) else IS_VALID_ATOL
             is_int = max_abs_diff < atol
             if self.bit_width >= 2:
                 if self.signed:

diff --git a/src/brevitas_examples/imagenet_classification/ptq/ptq_common.py b/src/brevitas_examples/imagenet_classification/ptq/ptq_common.py
@@ -441,9 +441,7 @@ def kwargs_prefix(prefix, weight_kwargs):
         if act_quant_type == 'asym' and act_quant_percentile is not None:
             act_quant = act_quant.let(**{'low_percentile_q': 100 - act_quant_percentile})
     if sym_act_quant is not None:
-        sym_act_quant = sym_act_quant.let(
-            **{
-                'high_percentile_q': act_quant_percentile})
+        sym_act_quant = sym_act_quant.let(**{'high_percentile_q': act_quant_percentile})
 
     weight_quant_dict = {'weight_quant': weight_quant}