NDT2EvalTrainer

class ts1.models.pretrained.NDT2EvalTrainer(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 (NDT2) – The instantiated model to link.

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