MtMEvalTrainer

class ts2.models.pretrained.MtMEvalTrainer(cfg, rank, world_size)

Bases: ts2.ts2_eval_trainer.TS2EvalTrainer

TS2 trainer for MtM, prompting the forward with the task’s mask-mode token.

Training corrupts the input with the configured masker and scores the loss on the corrupted entries; inference always prompts with TASK_MASK_MODE[task].

predict(X, mask=None, mask_timestamps=None, mask_units=None)

Generate model predictions for a batch.

Override this if the model requires masking information at inference time (e.g. for masked autoencoder models).

Parameters:

• X – Input batch dict containing model_inputs.

• mask – Optional boolean mask for masked prediction.

• mask_timestamps – Optional timestamps for masked positions.

• mask_units – Optional unit indices for masked positions.

Returns:

Model output predictions.

Return type:

torch.Tensor

setup_model(ckpt=None)

Instantiate and register the model.

Instantiates the model from the Hydra config and registers it for checkpointing. Override this to load pretrained weights or modify the model architecture.

Parameters:

ckpt (Optional[dict]) – Checkpoint dictionary, available for loading pretrained weights.

link_model(model, ckpt)

Link datasets to the model.

Attaches the configured transforms and the model input_fn to the train/val/test dataset transform pipelines, then calls model.link_datasets to register the datasets with the model.

Parameters:

• model (BaseModel) – The instantiated model to link.

• ckpt (dict | None) – Checkpoint dictionary, passed through for subclass use.

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 (dict) – Input batch dict containing model_inputs.

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

Returns:

Scalar loss.

Return type:

torch.Tensor