Scalar torch Tensors not recognized by (NoHeader)TensorSerializer -> Serialized as Pickle

[enrico-stauss]

🐛 Bug

I noticed that some scalar valued torch tensors in my dataset are being inferred as Pickle datatype. A quick investigation revealed, that neither the TensorSerialize nor the NoHeaderTensorSerializer can_serialize scalar valued tensors due to the conditions and len(item.shape) == 1 and and len(item.shape) > 1 respectively.

This is a quick fix and I'm happy to submit a PR. Am I correct to assume that a scalar valued tensor would best be serialized by the TensorSerializer? I think I recall that the NoHeaderTensorSerializer was designed specifically for LLM training use-cases?

To Reproduce

Execute the following snippet and you'll see a console printout indicating that the ``['pickle'] data format is inferred.

Code sample

import torch
from litdata import optimize


def fn(i):
    return torch.tensor(i, dtype=torch.float32)


if __name__ == "__main__":
    optimize(
        fn=fn,
        inputs=list(range(1000)),
        output_dir="litdata/test",
        chunk_bytes=f"{1}Mb",
    )

Expected behavior

Additional context

Environment detail

• PyTorch Version (e.g., 1.0):
• OS (e.g., Linux):
• How you installed PyTorch (conda, pip, source):
• Build command you used (if compiling from source):
• Python version:
• CUDA/cuDNN version:
• GPU models and configuration:
• Any other relevant information:

[tchaton]

Hey @enrico-stauss. Good catch. Please, submit a PR :)