This commit is contained in:
sayakpaul
2024-12-27 15:02:21 +05:30
parent 4461cdd98f
commit 84c17561eb
2 changed files with 8 additions and 1 deletions
+2 -1
View File
@@ -47,7 +47,7 @@ from .utils.data_utils import should_perform_precomputation
from .utils.file_utils import string_to_filename
from .utils.optimizer_utils import get_optimizer
from .utils.memory_utils import get_memory_statistics, free_memory, make_contiguous
from .utils.torch_utils import unwrap_model, align_device_and_dtype
from .utils.torch_utils import unwrap_model, align_device_and_dtype, expand_tensor_to_dims
from .utils.checkpointing import get_latest_ckpt_path_to_resume_from, get_intermediate_ckpt_path
@@ -696,6 +696,7 @@ class Trainer:
device=accelerator.device,
dtype=weight_dtype,
)
sigmas = expand_tensor_to_dims(sigmas, ndim=latent_conditions["latents"].ndim)
noisy_latents = (1.0 - sigmas) * latent_conditions["latents"] + sigmas * noise
latent_conditions.update({"noisy_latents": noisy_latents})