Image Generation while training?

[nicofirst1]

I've been trying to the DALLE model generate images while on training mode.
The generate_images method has the @torch.no_grad() and @eval_decorator decorators that do not allow training.
My solution was to copy-paste the same method and remove those decorators but I'm not sure it is the best solution.

My intent is to generate an image and feed it to another architecture and I would like the gradient to flow through the entire process.
Does anyone know how to do this?