PredictionsWriter#

class ibl_bwb_eval.predictions.PredictionsWriter(enable, base_path, task, seed, label=None, rank=0, metadata=None, path_fn=None)[source]#

Bases: object

Buffer tensors from an eval loop and save them as a .safetensors file.

The output path is structured as:

base_path / label / task / [path_fn(metadata)] / seed_{seed}.safetensors

path_fn is a callable that receives the accumulated path and metadata, and returns a middle segment of the path. This lets each task suite define its own sub-hierarchy (e.g. a recording id). metadata is also embedded in the file header alongside the common fields (label, task, seed) and ibl_bwb_eval_version, ibl_bwb_eval.__version__ at write time, which traces a submission back to the writer that produced it. Read that one back with version_of().

Fields are recorded in one of two ways:

  • add() appends a per-batch chunk; repeated calls for the same key are concatenated along dim 0 at save time.

  • set() stores a single value once, for a field that is constant across batches or already aggregated.

All values must be torch.Tensor. They are detached and moved to CPU on ingest. Floating-point tensors are cast to float32 by default; pass dtype= to add()/set() to use a lower-precision dtype instead (e.g. torch.float16 for bulky prediction tensors), avoid this for fields used for exact alignment (e.g. timestamps), since low precision can collapse distinct values. save() writes on rank 0 only and is a no-op if nothing was recorded.

On this page