fix: Replaced FloatTensor,IntTensor w/ Tensor

Signed-off-by: Brandon Groth <[email protected]>

Author: BrandonGroth

fms_mo/quant/quantizers.py

@@ -2544,7 +2544,7 @@ def asymmetric_linear_quantization_params(
         return scale, zero_point
 
 
-def clamp(input_tensor: torch.FloatTensor, clamp_min, clamp_max, inplace=False):
+def clamp(input_tensor: torch.Tensor, clamp_min, clamp_max, inplace=False):
     """
     Returns:
         Clamped Torch Tensor.

fms_mo/quant_refactor/base_quant.py

@@ -175,7 +175,7 @@ class Quantizer(torch.nn.Module):
 
     def __init__(
         self,
-        num_bits: torch.IntTensor,
+        num_bits: torch.Tensor,
         dequantize: bool = True,
         qscheme: Qscheme = torch.per_tensor_symmetric,
         use_PT_native_Qfunc: bool = False,
@@ -189,7 +189,7 @@ def __init__(
         Init Quantizer Class
 
         Args:
-            num_bits (torch.IntTensor): Number of bit for quantization.
+            num_bits (torch.Tensor): Number of bit for quantization.
             dequantize (bool, optional): Return dequantized or int tensor. Defaults to True.
             qscheme (Qscheme, optional): Quantization scheme.
                 Defaults to Qscheme(unit="perT", symmetric=True).
@@ -257,7 +257,7 @@ def set_quantizer(self):
             f"Quantizer selection is not implemented for quantizer {self}"
         )
 
-    def forward(self, input_tensor: torch.FloatTensor):
+    def forward(self, input_tensor: torch.Tensor):
         """
         General forward() function for quantizer classes.
 
@@ -266,13 +266,13 @@ def forward(self, input_tensor: torch.FloatTensor):
             To STE functions without calling calc_qparams()
 
         Args:
-            input_tensor (torch.FloatTensor): Tensor to be quantized.
+            input_tensor (torch.Tensor): Tensor to be quantized.
 
         Raises:
             ValueError: Single-sided qscheme has tensor min < 0.0
 
         Returns:
-            torch.FloatTensor: Dequantized or Quantized output tensor.
+            torch.Tensor: Dequantized or Quantized output tensor.
         """
         if self.qscheme.single_sided and input_tensor.min() < 0.0:
             raise ValueError(

fms_mo/quant_refactor/get_quantizer_new.py

@@ -41,8 +41,8 @@
 def get_activation_quantizer_new(
     qa_mode: str = "PACT",
     nbits: int = 32,
-    clip_val: torch.FloatTensor = None,
-    clip_valn: torch.FloatTensor = None,
+    clip_val: torch.Tensor = None,
+    clip_valn: torch.Tensor = None,
     non_neg: bool = False,
     align_zero: bool = True,  # pylint: disable=unused-argument
     extend_act_range: bool = False,
@@ -215,8 +215,8 @@ def get_activation_quantizer_new(
 def get_weight_quantizer_new(
     qw_mode: str = "SAWB+",
     nbits: int = 32,
-    clip_val: torch.FloatTensor = None,
-    clip_valn: torch.FloatTensor = None,
+    clip_val: torch.Tensor = None,
+    clip_valn: torch.Tensor = None,
     align_zero: bool = True,
     w_shape: torch.Size = None,
     recompute: bool = False,  # pylint: disable=unused-argument