Fix using .lstrip() in hivemind.compression (#578)

learning-at-home · Jul 21, 2023 · da130cd · da130cd
1 parent ec1d7fe
commit da130cd
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/hivemind/compression/base.py b/hivemind/compression/base.py
@@ -84,7 +84,7 @@ class NoCompression(CompressionBase):
     def compress(self, tensor: torch.Tensor, info: CompressionInfo, allow_inplace: bool = False) -> runtime_pb2.Tensor:
         tensor = tensor.detach()
         shape = tensor.shape
-        dtype_name = str(tensor.dtype).lstrip("torch.")
+        dtype_name = str(tensor.dtype).replace("torch.", "")
         raw_data = tensor
         if tensor.dtype == torch.bfloat16:
             if USE_LEGACY_BFLOAT16:  # legacy mode: convert to fp32

diff --git a/hivemind/compression/quantization.py b/hivemind/compression/quantization.py
@@ -140,7 +140,7 @@ def quantize(
 
     def compress(self, tensor: torch.Tensor, info: CompressionInfo, allow_inplace: bool = False) -> runtime_pb2.Tensor:
         tensor = tensor.detach()
-        dtype_name = str(tensor.dtype).lstrip("torch.")
+        dtype_name = str(tensor.dtype).replace("torch.", "")
         if tensor.dtype == torch.bfloat16:
             tensor = tensor.to(torch.float32)