def METHOD_NAME(model, loader, criterion, optimizer, max_steps, grad_norm=None,
               edge_dropout=0.0):
    model.train()
    total_loss = total_examples = 0
    for i, (batch, batch_size, *args) in enumerate(loader):
        x = batch.x.to(model.device)
        adj_t = batch.adj_t.to(model.device)
        y = batch.y[:batch_size].to(model.device)
        train_mask = batch.train_mask[:batch_size].to(model.device)
        if train_mask.sum() == 0:
            continue
        # We make use of edge dropout on ogbn-products to avoid overfitting.
        adj_t = dropout(adj_t, p=edge_dropout)
        optimizer.zero_grad()
        # cut the args with only n_id, which makes the call to be _async=False
        args = args[:1]
        out = model(x, adj_t, batch_size, *args)
        loss = criterion(out[train_mask], y[train_mask])
        loss.backward()
        if grad_norm is not None:
            torch.nn.utils.clip_grad_norm_(model.parameters(), grad_norm)
        optimizer.step()
        total_loss += float(loss) * int(train_mask.sum())
        total_examples += int(train_mask.sum())
        # We may abort after a fixed number of steps to refresh histories...
        if (i + 1) >= max_steps and (i + 1) < len(loader):
            break
    return total_loss / total_examples