I have a dataset with 100 images which occupy around 120 MB and their masks occupy around 4.2 MB. I have written a dataloader. When I try to plot some examples using the following code, the process is killed because of running out of memory.
import torchvision
import matplotlib.pyplot as plt
from augmentations import * 


path = '/path_to_data_set'
batch_size = 1
augs = Compose([RandomRotate(10),
						RandomHorizontallyFlip()])
root = SegDataset(root=path, split='train', is_transform=True, augmentations = augs)
trainloader = data.DataLoader(root, batch_size=batch_size, num_workers=0)
for i, data in enumerate(trainloader):
	print(i)
	imgs, masks = data
	imgs = imgs.numpy()[:, ::-1, :, :]
	imgs = np.transpose(imgs, (0,2,3,1))
	f, axarr = plt.subplots(batch_size,2)
	for j in range(batch_size):
		axarr[j][0].imshow(imgs[j])
		axarr[j][1].imshow(dst.decode_segmap(labels.numpy()[j]))
	plt.show()
	a = input()
	if a =='ex':
		break
	else:
		plt.close()

I have tried num_workers=0 and other numbers , still same behaviour.