Fix shape error

Signed-off-by: Kyunggeun Lee <quic_kyunggeu@quicinc.com>
quic · Oct 31, 2023 · 870efda · 870efda
1 parent 6151d66
commit 870efda
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/TrainingExtensions/torch/src/python/aimet_torch/tensor_quantizer.py b/TrainingExtensions/torch/src/python/aimet_torch/tensor_quantizer.py
@@ -814,6 +814,8 @@ def _compute_updated_encoding(self) -> Union[libpymo.TfEncoding, List[libpymo.Tf
         scale, offset = self.compute_scaling_offset(encoding_min.float(), encoding_max.float())
         assert scale is not None
         assert offset is not None
+        scale = scale.expand_as(encoding_min)
+        offset = offset.expand_as(encoding_min)
 
         if not self.use_symmetric_encodings or self.is_unsigned_symmetric:
             # Calculate 'min' and 'max' based on 'delta' and 'offset'