为什么多卡训练DM的时候embed将显存都集中在了主卡上

[slyang2021]

` if 'BN' not in args.model: # for ConvNet loss = torch.tensor(0.0).to(args.device) for c in range(num_classes): img_real = get_images(c, args.batch_real) img_syn = image_syn[cargs.ipc:(c+1)args.ipc].reshape((args.ipc, channel, im_size[0], im_size[1]))

                if args.dsa:
                    seed = int(time.time() * 1000) % 100000
                    img_real = DiffAugment(img_real, args.dsa_strategy, seed=seed, param=args.dsa_param)
                    img_syn = DiffAugment(img_syn, args.dsa_strategy, seed=seed, param=args.dsa_param)

                output_real = embed(img_real).detach()
                output_syn = embed(img_syn)

                loss += torch.sum((torch.mean(output_real, dim=0) - torch.mean(output_syn, dim=0))**2)`

[PatrickZH]

Hi, we have used the following code to make computation distributed on multiple GPUs: https://github.com/VICO-UoE/DatasetCondensation/blob/a84efcc0636b6578e398bb5614edc506e0f513a0/main_DM.py#L144 We didn't find this problem in our experiments.