Skip to content

Commit

Permalink
[Converter] read variable by chunk if size > max value of int (#1636)
Browse files Browse the repository at this point in the history
* read by chunk if size > max value of int

* fix black
  • Loading branch information
minhthuc2502 authored Mar 6, 2024
1 parent cec65f1 commit b4b3ac0
Showing 1 changed file with 11 additions and 1 deletion.
12 changes: 11 additions & 1 deletion python/ctranslate2/specs/model_spec.py
Original file line number Diff line number Diff line change
Expand Up @@ -730,7 +730,17 @@ def num_bytes(self) -> int:
return self.tensor.numel() * self.tensor.element_size()

def to_bytes(self) -> bytes:
return ctypes.string_at(self.tensor.data_ptr(), self.num_bytes())
max_size = 2**31 - 1
num_bytes = self.num_bytes()
output = b""
offset = 0
while num_bytes > 0:
chunk_size = max_size if num_bytes > max_size else num_bytes
chunk = ctypes.string_at(self.tensor.data_ptr() + offset, chunk_size)
output += chunk
offset += chunk_size
num_bytes -= chunk_size
return output

def _to(self, dtype: str) -> Variable:
dtype = getattr(torch, dtype)
Expand Down

0 comments on commit b4b3ac0

Please sign in to comment.