Hi @ptrblck,
Thank you so much for your reply! If I only care about

time to finish one epoch
time of entire mini-batch
time of loading each mini-batch

then I don’t need to use torch.cuda.synchronize(), and instead I can just modify the code as
for epoch in range(n_epoch):
    epoch_tic = time.time()
    tic = time.time()
    for x_batch, y_batch in training_loader:
        
        batch_load_time = time.time() - tic
        
        x_batch = x_batch.to(device)
        y_batch = y_batch.to(device)
        
        nnet = nnet.train()
        yhat = nnet(x_batch)
        loss = loss_fn(yhat, y_batch)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        batch_time = time.time() - tic
        tic = time.time()
    
    epoch_time = time.time() - epoch_tic

And correspondingly,


epoch_time is the time to finish one epoch

batch_load_time is the time to load one mini-batch of data

batch_time is the time to finish the training of one mini-batch

Is this correct?