Hi @ptrblck,
Thank you for your reply!
So the correct way should be adding torch.cuda.synchronize() in front of every time.time() call right?
for epoch in range(n_epoch):
    torch.cuda.synchronize()
    epoch_tic = time.time()
    tic = time.time()
    for x_batch, y_batch in training_loader:
        torch.cuda.synchronize()
        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()
        torch.cuda.synchronize()
        batch_time = time.time() - tic
        tic = time.time()
    torch.cuda.synchronize()
    epoch_time = time.time() - epoch_tic