update

Author: dsikka

src/compressed_tensors/quantization/lifecycle/forward.py

@@ -291,11 +291,11 @@ def _process_quantization(
             x = safe_permute(x, perm, dim=1)
 
         # Maintain all dimensions apart from the last dim, which is divided by the group_size
-        reshaped_dims = tuple(x.shape[:-1]) + (
+        reshaped_dims = (
             ceil(x.shape[-1] / group_size),
             group_size,
         )
-        x = torch.reshape(x, reshaped_dims)
+        x = x.unflatten(-1, reshaped_dims)
 
         if do_quantize:
             output = _quantize(
@@ -318,11 +318,7 @@ def _process_quantization(
                 global_scale=global_scale,
             )
 
-        original_shaped_dims = tuple(output.shape[:-2]) + (
-            output.shape[-1] * output.shape[-2],
-        )
-        output = torch.reshape(output, original_shaped_dims)
-
+        output = output.flatten(start_dim=-2)
         output = output.to(output_dtype)
 
         if not is_column_order:

src/compressed_tensors/quantization/utils/helpers.py

@@ -176,16 +176,14 @@ def compute_dynamic_scales_and_zp(
         QuantizationStrategy.GROUP,
     ):
 
-        reduce_dims = tuple(
-            idx for idx in range(len(value.shape) + 1) if idx not in range(value.dim())
-        )
+        reduce_dims = -1
         keep_dims = False
 
-        reshaped_dims = tuple(value.shape[:-1]) + (
+        reshaped_dims = (
             math.ceil(value.shape[-1] / args.group_size),
             args.group_size,
         )
-        value = torch.reshape(value, reshaped_dims)
+        value = value.unflatten(-1, reshaped_dims)
 
     else:
         supported_strategies = (