nvalchemi.training.LossWeightSchedule

class nvalchemi.training.LossWeightSchedule(*args, **kwargs)

Runtime-checkable protocol for loss-weight schedules.

Any object callable with signature (step: int, epoch: int) -> float, exposing a per_epoch attribute, and returning a rebuild recipe from to_spec() satisfies this protocol. Such objects are accepted inside ComposedLossFunction’s weights sequence or as the right-hand side of schedule * leaf. Concrete Pydantic schedules live in schedules.

per_epoch

If True, the schedule should advance by epoch instead of by step. This aligns loss-weight updates with training loops that update learning-rate schedules once per epoch.

Type:

bool

Parameters:

• step – Current global training step (0-indexed).

• epoch – Current epoch number (0-indexed).

Returns:

Scalar weight to apply to the associated loss term.

Return type:

float

to_spec()

Return a serializable spec that rebuilds this schedule.

Return type:

BaseSpec