1 files changed, 5 insertions, 2 deletions

diff --git a/training/functional.py b/training/functional.py
index 43b03ac..546aaff 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -2,6 +2,7 @@ from dataclasses import dataclass
 import math
 from contextlib import _GeneratorContextManager, nullcontext
 from typing import Callable, Any, Tuple, Union, Optional, Protocol
+from types import MethodType
 from functools import partial
 from pathlib import Path
 import itertools
@@ -108,6 +109,7 @@ def save_samples(
     output_dir: Path,
     seed: int,
     step: int,
+    validation_prompts: list[str] = [],
     cycle: int = 1,
     batch_size: int = 1,
     num_batches: int = 1,
@@ -136,7 +138,6 @@ def save_samples(
 
     if val_dataloader is not None:
         datasets.append(("stable", val_dataloader, generator))
-        datasets.append(("val", val_dataloader, None))
 
     for pool, data, gen in datasets:
         all_samples = []
@@ -165,7 +166,6 @@ def save_samples(
                     guidance_scale=guidance_scale,
                     sag_scale=0,
                     num_inference_steps=num_steps,
-                    output_type=None,
                 ).images
 
                 all_samples.append(torch.from_numpy(samples))
@@ -803,4 +803,7 @@ def train(
     accelerator.unwrap_model(text_encoder, keep_fp32_wrapper=False)
     accelerator.unwrap_model(unet, keep_fp32_wrapper=False)
 
+    text_encoder.forward = MethodType(text_encoder.forward, text_encoder)
+    unet.forward = MethodType(unet.forward, unet)
+
     accelerator.free_memory()