fix typo (#241)

neuralmagic · Jan 20, 2025 · 9a0afd6 · 9a0afd6
1 parent 0ba31bf
commit 9a0afd6
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/src/compressed_tensors/compressors/sparse_compressors/sparse_24_bitmask.py b/src/compressed_tensors/compressors/sparse_compressors/sparse_24_bitmask.py
@@ -103,8 +103,10 @@ def from_compressed_data(
         :param bitmask: 2d bitmask of non-zero values
         :return: instantiated Sparse24BitMaskTensor
         """
-        if isinstance(shape, Tensor):
-            shape = shape.tolist()
+        if isinstance(shape, list):
+            shape = torch.tensor(shape)
+        if isinstance(shape, torch.Tensor):
+            shape = shape.flatten().tolist()
         return Sparse24BitMaskTensor(
             shape=shape, compressed=compressed, bitmask=bitmask
         )

diff --git a/src/compressed_tensors/utils/helpers.py b/src/compressed_tensors/utils/helpers.py
@@ -301,7 +301,7 @@ def pack_bitmasks(bytemasks: torch.Tensor) -> torch.Tensor:
 
 
 def unpack_bitmasks(
-    packed_bitmasks: torch.Tensor, original_shape: torch.Size
+    packed_bitmasks: torch.Tensor, original_shape: List[int]
 ) -> torch.Tensor:
     """
     Converts a bitmask tensor back to a bytemask tensor for use during decompression