diff options
Diffstat (limited to 'training')
| -rw-r--r-- | training/common.py | 2 |
1 files changed, 1 insertions, 1 deletions
diff --git a/training/common.py b/training/common.py index 67c2ab6..0b2ae44 100644 --- a/training/common.py +++ b/training/common.py | |||
| @@ -58,7 +58,7 @@ def generate_class_images( | |||
| 58 | torch.cuda.empty_cache() | 58 | torch.cuda.empty_cache() |
| 59 | 59 | ||
| 60 | 60 | ||
| 61 | def run_model( | 61 | def loss_step( |
| 62 | vae: AutoencoderKL, | 62 | vae: AutoencoderKL, |
| 63 | noise_scheduler: DDPMScheduler, | 63 | noise_scheduler: DDPMScheduler, |
| 64 | unet: UNet2DConditionModel, | 64 | unet: UNet2DConditionModel, |