yuanc:

loss = nn.L1Loss()(Y_pred, Y_train_tensor)


This line presumably needs to be
loss = nn.L1Loss()(Y_batch_pred, Y_batch)

Can you check if this helps? I cannot see anything else immediately that would cause a memory issue.
I also assume that the 16 in range(16) is handcrafted to fit your dataset. But it seems odd, given that you have 200k samples. What’s your batch_size? Because currently you go only through 16*batch_size samples. Maybe just let the inner for loop run indefinitely and break if len(X_batch) == 0.