compatibility fix

This commit is contained in:
Christian Zhou-Zheng 2024-06-10 14:00:13 -04:00
parent 854bd64a5d
commit 05b183fe7b

View file

@ -346,10 +346,25 @@ class GGUFWriter:
if self.endianess == GGUFEndian.BIG:
tensor.byteswap(inplace=True)
for fout in self.fout:
self.write_padding(fout, fout.tell())
tensor.tofile(fout)
self.write_padding(fout, tensor.nbytes)
file_id = -1
for i, tensors in enumerate(self.tensors):
if len(tensors) > 0:
file_id = i
break
fout = self.fout[file_id]
# pop the first tensor info
# TODO: cleaner way to get the first key
first_tensor_name = [name for name, _ in zip(self.tensors[file_id].keys(), range(1))][0]
ti = self.tensors[file_id].pop(first_tensor_name)
assert len(ti.shape) == len(tensor.shape)
assert all(dim1 == dim2 for dim1, dim2 in zip(ti.shape, tensor.shape))
assert ti.nbytes == tensor.nbytes
self.write_padding(fout, fout.tell())
tensor.tofile(fout)
self.write_padding(fout, tensor.nbytes)
self.state = WriterState.WEIGHTS