Function validate

pretrain/main.py:642–669

Run validation and calculate loss.

(
    fabric: lightning.Fabric,
    model: torch.nn.Module,
    val_dataloader: DataLoader,
    eval_iters: int,
)

@torch.no_grad()
def validate(
    fabric: lightning.Fabric,
    model: torch.nn.Module,
    val_dataloader: DataLoader,
    eval_iters: int,
) -> torch.Tensor:
    """Run validation and calculate loss."""

    fabric.print('Validating ...')
    model.eval()

    losses = torch.zeros(eval_iters, device=fabric.device)
    for k, val_data in enumerate(val_dataloader):
        if k >= eval_iters:
            break
        input_ids = val_data[:, 0 : model.config.block_size].contiguous()
        targets = val_data[:, 1 : model.config.block_size + 1].contiguous()
        logits = model(input_ids)
        loss = chunked_cross_entropy(logits, targets, chunk_size=0)

        # loss_func = FusedCrossEntropyLoss()
        # loss = loss_func(logits, targets)
        losses[k] = loss.item()

    out = losses.mean()

    model.train()
    return out


def create_dataloader(