speed up latents nan replace

kohya-ss · Dec 20, 2023 · 8c7d05a · 8c7d05a
1 parent 0908c54
commit 8c7d05a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train_network.py b/train_network.py
@@ -767,7 +767,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]