From 8aaa1967bd3d3a9b4b44e97e5432d23f2101cf51 Mon Sep 17 00:00:00 2001
From: Kohya S <ykumeykume@gmail.com>
Date: Thu, 15 Aug 2024 22:07:23 +0900
Subject: [PATCH] fix encoding latents closes #1456

---
 flux_train_network.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/flux_train_network.py b/flux_train_network.py
index b8ea5622..daa65c85 100644
--- a/flux_train_network.py
+++ b/flux_train_network.py
@@ -238,8 +238,8 @@ class FluxNetworkTrainer(train_network.NetworkTrainer):
         return noise_scheduler
 
     def encode_images_to_latents(self, args, accelerator, vae, images):
-        return vae.encode(images).latent_dist.sample()
-
+        return vae.encode(images)
+    
     def shift_scale_latents(self, args, latents):
         return latents