If your goal is to print, for each epoch, the cumulative loss in order to analyze in TensorBoard/WANDB the performances of the model you might do something like this:
cumulative_loss = 0
with torch.set_grad_enabled(train):
    for batch_idx, (inputs, targets) in enumerate(data_loader):
      inputs, targets = inputs.to(device, non_blocking=True), targets.to(device, non_blocking=True) 
      outputs = net(inputs) 
      loss = cost_function(outputs, targets) 
      loss2 = cost_function(outputs, targets)
      if train:
        loss.backward() # Backward pass
        optimizer.step() # Update parameters
        optimizer.zero_grad() # Resets the gradients

      #cumulative_loss += loss.item()  very slow
      cumulative_loss += loss2   #should not need synch

    cumulative_loss = cumulative_loss.item() #synch once per epoch instead of once per batch
    logging_loss_epoch(cumulative_loss)

After some experiments to me it seems that the performances are much better.