Code changes for gpu

[karthik1coder]

Inorder to run on gpu,what changes must be done?Can I know if you have used zero stage in accelerator to run this code?

[danbochman]

The only thing you need to do to use GPU is to remove the accelerate_cpu parameter in the example script. After that, it is up to your accelerate config setup.

https://github.com/fashn-AI/tryondiffusion/blob/master/examples/test_tryon_imagen_trainer.py#L96

Regarding DeepSpeed - I haven't been able to integrate it for this repository. I am compiling the model with torch.compile and training with DDP.

[karthik1coder]

yeah i tried already removing accelerate_cpu the only problem is that i am facing oom error.Thank you