NEDSEvalTrainer

class ts1.models.pretrained.NEDSEvalTrainer(cfg, rank, world_size)

Bases: ts1.ts1_eval_trainer.TS1EvalTrainer

setup_model(ckpt)

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 (dict | None) – Checkpoint dictionary, available for loading pretrained weights.

link_model(model, ckpt)

Link datasets to the model and configure the readout head.

Attaches model-specific transforms and input_fn to the train/val/test dataset pipelines, then calls model.link_datasets and model.configure_readout to wire up the readout head for the current task.

Parameters:

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

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

predict(X, target_timestamps=None, target_mask=None)

Generate model predictions for a batch.

Override this if the model requires target timestamps or mask at inference time (e.g. for masked prediction models).

Parameters:

• X – Input batch dict containing model_inputs.

• target_timestamps – Optional target timestamps tensor.

• target_mask – Optional boolean mask tensor for the targets.

Returns:

Model output predictions.

Return type:

torch.Tensor