fix: check against NaNs in images

LuxDL · Jan 5, 2025 · fafb9b7 · fafb9b7
1 parent cbf0d8d
commit fafb9b7
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/examples/DDIM/main.jl b/examples/DDIM/main.jl
@@ -307,6 +307,10 @@ end
 
 function save_images(output_dir, images::Vector{<:AbstractMatrix{<:RGB}})
     for (i, img) in enumerate(images)
+        if any(isnan, img)
+            @warn "NaN image found in the generated images. Skipping..."
+            continue
+        end
         save(joinpath(output_dir, "img_$(i).png"), img)
     end
 end