nvalchemi.hooks.DynamicsContext

class nvalchemi.hooks.DynamicsContext(*, batch, model=None, global_rank=0, workflow=None, step_count=0, converged_mask=None)

Context object passed to dynamics hooks.

Parameters:

• batch (Batch | None)

• model (BaseModelMixin | None)

• global_rank (int)

• workflow (Any)

• step_count (int)

• converged_mask (torch.Tensor | None)

step_count

Current dynamics step number.

Type:

int

converged_mask

Boolean mask of samples that converged at the current hook stage. None when convergence has not fired for this dispatch.

Type:

torch.Tensor | None