Update eval.py
Browse files
eval.py
CHANGED
|
@@ -303,6 +303,8 @@ def train_loop(config, model, noise_scheduler, optimizer, train_dataloader, lr_s
|
|
| 303 |
model, optimizer, train_dataloader, lr_scheduler
|
| 304 |
)
|
| 305 |
|
|
|
|
|
|
|
| 306 |
if accelerator.is_main_process:
|
| 307 |
evaluate(config, 1, config.pipeline)
|
| 308 |
|
|
|
|
| 303 |
model, optimizer, train_dataloader, lr_scheduler
|
| 304 |
)
|
| 305 |
|
| 306 |
+
pipeline = DDPMPipeline(unet=accelerator.unwrap_model(model), scheduler=noise_scheduler)
|
| 307 |
+
|
| 308 |
if accelerator.is_main_process:
|
| 309 |
evaluate(config, 1, config.pipeline)
|
| 310 |
|