Diffusion model using pytorch memory issue

Hi

I am trying to replicate this example . Only my input data has batch size of 5,1,70,70. So I am not using the data used in this example. My question is about the inference section. When for one batch(size 5,1,70,70), if i do timesteps iterative addition, i get memory issues. In this example it is1000 timesteps. I am unable to run it for more than 50 time steps. How this code can be changed to make it more memory efficient to achieve this 1000 timestep iterations? Also, an explanation would be appreciated as why it is too memory intensive?

Thanks