Merge pull request #1009 from liubo0902/main

speed up latents nan replace
kohya-ss · Dec 21, 2023 · 0676f1a · 0676f1a
2 parents 6b7823d + 8c7d05a
commit 0676f1a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train_network.py b/train_network.py
@@ -750,7 +750,7 @@ def remove_model(old_ckpt_name):
                             # NaNが含まれていれば警告を表示し0に置き換える
                             if torch.any(torch.isnan(latents)):
                                 accelerator.print("NaN found in latents, replacing with zeros")
-                                latents = torch.where(torch.isnan(latents), torch.zeros_like(latents), latents)
+                                latents = torch.nan_to_num(latents, 0, out=latents)
                         latents = latents * self.vae_scale_factor
                     b_size = latents.shape[0]