Fix comments

crates/burn-jit/src/kernel/quantization/dequantize.rs

@@ -47,7 +47,7 @@ pub(crate) fn dequantize_per_tensor_affine_int8_kernel(
     #[comptime] scheme: QuantizationScheme,
 ) {
     // Last two positions contain the qparams
-    if ABSOLUTE_POS >= output.len() - 2 {
+    if ABSOLUTE_POS >= input.len() - 2 {
         return;
     }
 

crates/burn-jit/src/kernel/quantization/qtensor.rs

@@ -27,8 +27,8 @@ impl QParams {
             QuantizationScheme::PerTensorAffine(_) => match comptime!(tensor.line_size()) {
                 // For line size of 1, scale is the last value in the buffer
                 1 => f32::bitcast_from(tensor[len - 1][tensor.line_size() - 1]),
-                // For any other line size > 1, scale and zero-point offset are the first two elements of the last line
-                _ => f32::bitcast_from(tensor[len - 1][tensor.line_size() - 2]),
+                // For any other line size > 1, scale and zero-point offset are the last two elements
+                _ => f32::bitcast_from(tensor[len - 1][tensor.line_size() - 1]),
             },
             // Symmetric quantization only contains the scaling factor as the last element
             QuantizationScheme::PerTensorSymmetric(_) => {
@@ -40,13 +40,12 @@ impl QParams {
     /// Get the zero-point offset.
     pub fn offset(&self, tensor: &QTensor) -> i32 {
         let len = tensor.len();
-        let line_size = comptime!(tensor.line_size());
         match comptime!(self.scheme) {
-            QuantizationScheme::PerTensorAffine(_) => match line_size {
+            QuantizationScheme::PerTensorAffine(_) => match comptime!(tensor.line_size()) {
                 // For line size of 1, the zero-point offset is the penultimate value in the buffer
-                1 => i32::cast_from(tensor[len - 2][line_size]),
-                // For any other line size > 1, scale and zero-point offset are the first two elements of the last line
-                _ => i32::cast_from(tensor[len - 1][line_size]),
+                1 => i32::cast_from(tensor[len - 2][tensor.line_size() - 1]),
+                // For any other line size > 1, scale and zero-point offset are the last two elements
+                _ => i32::cast_from(tensor[len - 1][tensor.line_size() - 2]),
             },
             // Symmetric quantization only contains the scaling factor, so we return 0 for the zero-point offset
             QuantizationScheme::PerTensorSymmetric(_) => 0,