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.