import nvtx
from torch.cuda import profiler

def copy_data(batch):
    data, targets = batch
    data_gpu = data.to(DEVICE)
    targets_gpu = targets.to(DEVICE)
    return data_gpu, targets_gpu


def compute_step(model, batch, optimizer):
    data, targets = batch
    output = model(data)
    loss = torch.nn.functional.cross_entropy(output, targets)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    return loss


data_iter = iter(train_loader)

for i in range(TOTAL_STEPS):

    if i == WARMUP_STEPS:
        # start nsys profiler
        torch.cuda.synchronize()
        start_time = time.perf_counter()
        profiler.start()
    elif i == WARMUP_STEPS + PROFILE_STEPS:
        # stop nsys profiler
        torch.cuda.synchronize()
        profiler.stop()
        end_time = time.perf_counter()

    with nvtx.annotate(f"Batch {i}", color="blue"):
        with nvtx.annotate("get batch", color="red"):
            batch = next(data_iter)
        with nvtx.annotate("copy batch", color="yellow"):
            batch = copy_data(batch)
        with nvtx.annotate("Compute", color="green"):
            compute_step(model, batch, optimizer)

total_time = end_time - start_time
throughput = PROFILE_STEPS / total_time
print(f"Throughput: {throughput:.2f} steps/sec")