Hi,
I don’t know at all how does it internally replace the values.
You can try
        sample[:, :, 0].zero_()
        sample[:, :, 1].zero_()

(Assuming that sample is a torch tensor)
That’s in-place replacement which should consume no extra memory.
It’s normal a single iteration takes a long time since you are running it on CPU.
You can try to use google colab