NDTEvalTrainer

class ts2.models.single_session.NDTEvalTrainer(cfg, rank, world_size)

Bases: ts2.ts2_eval_trainer.TS2EvalTrainer

TS2 trainer for the single-session NDT.

Training corrupts the input with the configured masker and scores the loss on the corrupted entries.

setup(ckpt=None)

Set up all trainer components.

Called automatically by the base constructor. Initializes tasks, data loaders, model, DDP wrapping, optimizer, and finetuning strategy in order.

Parameters:

ckpt (Optional[dict]) – Checkpoint dictionary loaded from disk, or None for a fresh run.

training_step(X, y)

Compute loss for one training batch.

Override this to customize masking, auxiliary losses, or any model-specific training logic. The base implementation is a plain forward pass against the benchmark-fixed target.

Parameters:

• X – Input batch dict containing model_inputs.

• y – Target batch dict containing values (benchmark-fixed spike counts).

Returns:

Scalar loss.

Return type:

torch.Tensor