implement squeeze in int_quant_tensor

Xilinx · May 1, 2024 · 8283979 · 8283979
1 parent 5ae5719
commit 8283979
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/brevitas/quant_tensor/int_quant_tensor.py b/src/brevitas/quant_tensor/int_quant_tensor.py
@@ -235,6 +235,9 @@ def dim(self):
     def add(self, other):
         return self + other
 
+    def unsqueeze(self, *args, **kwargs):
+        return self.value.unsqueeze(*args, **kwargs)
+
     @staticmethod
     def cat(tensors, dim, out=None):
         if out is not None: