I have removed them, this is up-to-date ram-hogging code.
class SegDataset(data.Dataset):
	def __init__(self, root, split="train", is_transform = False):
		self.root=root
		self.split = split
		self.augmentations = augmentations
		self.is_transform = is_transform
		self.n_classes = 3
		self.files = collections.defaultdict(list)
		self.images_base = os.path.join(self.root, self.split)
		self.masks_base = os.path.join(self.root, 'masks' + '/' + self.split)

		for split in ['train', 'val', 'test']:
			file_list = recursive_glob(rootdir=self.root + '/' + split , suffix='.jpg')
			self.files[split] = file_list

	def __len__(self):
		return len(self.files[self.split])

	def __getitem__(self, index):
		img_path = self.files[self.split][index].rstrip()
		img_name = img_path.split(os.sep)[-1].split('.')[0]
		mask_name = img_name + 	'_mask.png'
		mask_path = self.masks_base + '/' + mask_name
		
		img = io.imread(img_path)
		mask = io.imread(mask_path)
		if not (len(img.shape) == 3 and len(mask.shape) == 2):
			return self.__getitem__(np.random.randint(0, self.__len__()))

		if self.is_transform:
			img, mask = self.transform(img, mask)

		return img, mask
    
	def transform(self, img, mask):
		img = img[:, :, ::-1]
		img = img.astype(np.float64)
		img = img.transpose(2, 0, 1)
		classes = np.unique(mask)
		mask = mask.astype(int)
		assert(np.all(classes == np.unique(mask)))
                labels = encode_mask(mask)
		img = torch.from_numpy(img).float()
		mask = torch.from_numpy(mask).long()
		

		return img, mask

	def get_colors(self):
		return np.asarray([[0, 0, 0],
						[125, 0, 0],
						[0, 125, 0]])

	def encode_mask(self, mask):
		mask = mask.astype(int)
		labels = np.zeros(mask.shape[0], mask.shape[1], dtype=np.int16)
		for i, label in enumerate(self.get_colors()):
			labels[np.where(np.all(mask == label, axis=-1))[:2]] = 1
		labels = labels.astype(int)
		return labels

The transform method returns only the mask not the encoded labels and I set the is_transform to False. When I run it the process gets killed after some time. If I remove the transform then the code works as expected.