Fix (proxy/runtime): Typo fix on type hint

Xilinx · Sep 3, 2024 · 16847f4 · 16847f4
1 parent 71c93c8
commit 16847f4
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/brevitas/proxy/runtime_quant.py b/src/brevitas/proxy/runtime_quant.py
@@ -206,7 +206,9 @@ def bit_width(self, force_eval=True):
         return self.retrieve_attribute('bit_width', force_eval)
 
     def create_quant_tensor(
-            self, qt_args: Tuple[Any], x: Optional[IntQuantTensor] = None) -> IntQuantTensor:
+            self,
+            qt_args: Union[Tensor, Tuple[Any]],
+            x: Optional[IntQuantTensor] = None) -> IntQuantTensor:
 
         if x is None:
             out = IntQuantTensor(*qt_args, self.is_signed, self.training)