Fix weight residual computation

Xilinx · Dec 18, 2024 · 33bd3f7 · 33bd3f7
1 parent e7ef917
commit 33bd3f7
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/brevitas/quant_tensor/groupwise_float_quant_tensor.py b/src/brevitas/quant_tensor/groupwise_float_quant_tensor.py
@@ -109,7 +109,7 @@ def expand(self):
         # Then, we unbind the tensor along the group_dim shape, and drop the padded columns
         # Finally, we stack the remaining tensors
         unpadding_shape = final_shape[self.group_dim]
-        residual = curr_shape[self.group_dim] - unpadding_shape
+        residual = new_value.shape[self.group_dim] - unpadding_shape
 
         if residual > 0:
             new_value = torch.stack(

diff --git a/src/brevitas/quant_tensor/groupwise_int_quant_tensor.py b/src/brevitas/quant_tensor/groupwise_int_quant_tensor.py
@@ -95,7 +95,7 @@ def expand(self):
         # Then, we unbind the tensor along the group_dim shape, and drop the padded columns
         # Finally, we stack the remaining tensors
         unpadding_shape = final_shape[self.group_dim]
-        residual = curr_shape[self.group_dim] - unpadding_shape
+        residual = new_value.shape[self.group_dim] - unpadding_shape
 
         if residual > 0:
             new_value = torch.stack(