Use multi-GPU to calc FID score

[DQSSSSS]

I'm training the diffusion model and I want to calculate the FID score for each some steps.
Because the FID calculation needs to generate so many samples, e.g. 5000, I hope to distribute them to each GPU.
I know I can create a Callback class and implement it in on_validation_epoch_start: generate 5000 samples, get the result, and log it, but it's on a single GPU.

[KasuganoLove]

Have you found a solution to this problem? I currently have an issue with how to test FID using PyTorch Lightning. Taking the CelebA dataset as an example, if I write a dataloader to iterate through the entire dataset as the validation set, how can we control the number of generated fake images for each val_step?

[KasuganoLove]

here is the code to be completed:

from torchmetrics.image.fid import FrechetInceptionDistance
self.val_FID = FrechetInceptionDistance(normalize=True).set_dtype(torch.float64)

    def validation_step(self, batch, batch_idx: int):
        targets = batch
        # initialize X_0
        X_0 = targets
        # calculate FID
        self.val_FID.update(X_0, real=True)
        # how to generate 50K fake images, while the val real images number is ~20k?
        ....
        ....
        ....
        self.log('val_FID', self.val_FID, on_epoch=True, prog_bar=True, logger=True, sync_dist=True)