I tried the approach but this time it didn’t fit the memory. I guess that PyTorch accumulates all batch information when I do this way and that doesn’t have the possibility of fitting into memory. Does anybody have a better idea to accomplish this? Thanks.
    scheduler.step()
    optimizer.zero_grad()
    losses = torch.zeros(len(train_loader),).float().cuda()
    for i, (batch_features, _, _, batch_targets) in enumerate(train_loader):

        ...
        outputs = model(features).float().cuda()

        batch_loss = custom_loss(outputs, batch_targets)

        losses[i] = batch_loss
    entro = losses / losses.sum()
    entro = (entro * torch.log(entro)).mean()
    loss = losses.mean() * entro * -1.0
    loss.backward()
    optimizer.step()