crowsonkb / k-diffusion

Karras et al. (2022) diffusion models for PyTorch
MIT License
2.21k stars 371 forks source link

Add transform_last for BrownianTreeNoiseSampler #68

Open Disty0 opened 1 year ago

Disty0 commented 1 year ago

torch.Generator is not supported on Intel so we have to run BrownianTreeNoiseSampler on the CPU with transform and send results back to the GPU with transform_last.

It will throw this error if we try sending it to the GPU outside of the class: AttributeError: 'BrownianTreeNoiseSampler' object has no attribute 'to'

Example use: return BrownianTreeNoiseSampler(x.to("cpu"), sigma_min, sigma_max, seed=current_iter_seeds, transform=lambda x: x.to("cpu"), transform_last=lambda x: x.to("xpu"))