def METHOD_NAME(model, device, train_loader, optimizer, epoch, log_interval, dry_run):
    model.METHOD_NAME()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % log_interval == 0:
            wandb.log(
                {"epoch": epoch, "step": batch_idx * len(data), "loss": loss.item()}
            )
            if dry_run:
                break