the grad calculation takes up a lot of memory

[PanXiebit]

Smooth-Diffusion/train_smooth_diffusion.py

Line 312 in 5522761

grad, = autograd.grad(

grad, = autograd.grad(
        outputs=(fake_img * noise).sum(), inputs=latents, create_graph=True
    )

The calculation of gradients is memory inefficient and lacks support for flast-attention. Consequently, when training with the reg_loss, it becomes necessary to reduce the batch_size.

[Pakase]

Hi, I wander how much the extra memory cost may be? Could you give me a rough number?