aljaz/DALLE2-pytorch
1
1
Fork 0
mirror of https://github.com/lucidrains/DALLE2-pytorch.git synced 2025-12-19 17:54:20 +01:00
Code Issues Packages Projects Releases Wiki Activity

extra insurance that diffusion prior is on the correct device, when using trainer with accelerator or device was given

Browse Source
This commit is contained in:
Phil Wang
2022-07-07 10:08:33 -07:00
parent 46be8c32d3
commit ec7cab01d9
2 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
dalle2_pytorch/trainer.py
View File

@@ -192,6 +192,7 @@ class DiffusionPriorTrainer(nn.Module):
self.device = diffusion_prior_device self.device = diffusion_prior_device
else: else:
self.device = accelerator.device if exists(accelerator) else device self.device = accelerator.device if exists(accelerator) else device
diffusion_prior.to(self.device)
# save model # save model

2
dalle2_pytorch/version.py
View File

@@ -1 +1 @@
__version__ = '0.16.16' __version__ = '0.16.17'