Out of CUDA mem when training with streaming dataset

I intend to train my model in a subset of WebSight dataset. However, when i adjust subset_size into bigger size, it causes the out of CUDA memory error, though I remain the same batch size of 2.

dataset = load_dataset("HuggingFaceM4/WebSight", "v0.2", split='train', streaming=True)
subset_size = 1000
train_rate = 0.8
sub_ds = dataset.take(subset_size)
train_ds = sub_ds.take(round(subset_size * train_rate))
eval_ds = sub_ds.skip(round(subset_size * train_rate))

Maybe this will help you, I’m also new

1 Like

Can we see the training code?