In your code you are accumulating stats in:
        train_accuracies += im_acc
        train_ious += im_iou
        train_mses += pc_mse

which could increase the memory usage, if some of these tensors are still attached to the computation graph since the entire graph would also be stored in each iteration.
Assuming you want to track these statistics without calling backward() on any of these tensors, make sure to .detach() the tensors before adding them or call .item() in case it’s a scalar value.