I have a theory: Maybe those two lines
seg_loss  = ce_loss_criterion(seg_out, seg_gt_masked_1) * weight_1
seg_loss += ce_loss_criterion(seg_out, seg_gt_masked_2) * weight_2

in combination with
seg_loss.backward()
lead to the situation were the network is implicitly split into two copies and each loss is backpropagated separately. Then it is no surprise that I observe the consumption of the double amount of memory once in a while. Can anyone confirm or disprove this?