Contrastive Guidace cannot work on multiple GPUs

[LQF39466]

I am able to run standard DiffPure on multiple GPUs and achieve satisfactory robust accuracy. However, when I try to execute the contrastive guided version with the same GPU configuration, I encounter an error where self.optimizer in the reset_drift_counter() function receives an empty parameter list. Changing my GPU configuration to use only cuda:0 resolved the issue, suggesting that the contrastive guided diffusion in this repository may only operate on a single GPU, specifically the first one in the system.

Further investigation revealed that the embed_net in the ContrastiveGuidedDiffusion class seems to be fixed to cuda:0, regardless of any modifications to the GPU configuration.

I would like to understand whether this is a bug or an inherent limitation of the contrastive guidance implementation.

[tenghuilee]

Hello. Thank you so much for reporting this issue and providing such a detailed description. We only run the code on single GPU per task. We did not try on multi-gpu situations. Thank you again, we will try to fix this bug.