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

use eval vs training mode to determine whether to call backprop on trainer forward

Browse Source
This commit is contained in:
Phil Wang
2022-05-15 14:20:59 -07:00
parent 156fe5ed9f
commit 7b7a62044a
2 changed files with 7 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
dalle2_pytorch/train.py
View File

@@ -279,6 +279,8 @@ class DiffusionPriorTrainer(nn.Module):
loss = loss * chunk_size_frac loss = loss * chunk_size_frac
total_loss += loss.item() total_loss += loss.item()
if self.training:
self.scaler.scale(loss).backward() self.scaler.scale(loss).backward()
return total_loss return total_loss
@@ -406,6 +408,8 @@ class DecoderTrainer(nn.Module):
loss = loss * chunk_size_frac loss = loss * chunk_size_frac
total_loss += loss.item() total_loss += loss.item()
if self.training:
self.scale(loss, unet_number = unet_number).backward() self.scale(loss, unet_number = unet_number).backward()
return total_loss return total_loss

2
setup.py
View File

@@ -10,7 +10,7 @@ setup(
'dream = dalle2_pytorch.cli:dream' 'dream = dalle2_pytorch.cli:dream'
], ],
}, },
version = '0.2.31', version = '0.2.32',
license='MIT', license='MIT',
description = 'DALL-E 2', description = 'DALL-E 2',
author = 'Phil Wang', author = 'Phil Wang',