fix: weights must end with .weights.h5

ghaiszaher · Aug 25, 2024 · 71a2b60 · 71a2b60
1 parent c4048b1
commit 71a2b60
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/lib/train.py b/lib/train.py
@@ -134,10 +134,10 @@ def rgb_ratio_loss(self, clear_image, generated_image):
 
     def get_models_and_paths(self):
         import os
-        generator_clear2fog_weights_path = os.path.join(self.weights_path, 'generator_clear2fog.h5')
-        generator_fog2clear_weights_path = os.path.join(self.weights_path, 'generator_fog2clear.h5')
-        discriminator_clear_weights_path = os.path.join(self.weights_path, 'discriminator_clear.h5')
-        discriminator_fog_weights_path = os.path.join(self.weights_path, 'discriminator_fog.h5')
+        generator_clear2fog_weights_path = os.path.join(self.weights_path, 'generator_clear2fog.weights.h5')
+        generator_fog2clear_weights_path = os.path.join(self.weights_path, 'generator_fog2clear.weights.h5')
+        discriminator_clear_weights_path = os.path.join(self.weights_path, 'discriminator_clear.weights.h5')
+        discriminator_fog_weights_path = os.path.join(self.weights_path, 'discriminator_fog.weights.h5')
         models = [self.generator_clear2fog,
                   self.generator_fog2clear,
                   self.discriminator_clear,