multi-gpu training support?

[nanlliu]

it looks like it only supports a single GPU training rn. Do you have any plans to add multi-gpu training option any time soon? much appreciated

[wangqiang9]

use torch.nn.DataParallel is enough

[rlangefe]

Where exactly do you suggest putting that call? I tried it on both the GaussianDiffusion and Unet models, but it fails on both because it tries to access attributes that no longer exist when you call torch.nn.DataParallel on it.

[CerebralSeed]

Open the diffusion file, and run a search and replace model. ---> model.module.