diff --git a/src/mflux/dreambooth/lora_layers/lora_layers.py b/src/mflux/dreambooth/lora_layers/lora_layers.py
index 2ad62d4..edaa5cb 100644
--- a/src/mflux/dreambooth/lora_layers/lora_layers.py
+++ b/src/mflux/dreambooth/lora_layers/lora_layers.py
@@ -58,7 +58,7 @@ class LoRALayers:
             lora_layers = {**transformer_lora_layers, **single_transformer_lora_layers}
 
             weights = WeightHandler(
-                meta_data=MetaData(is_mflux=True),
+                meta_data=MetaData(mflux_version=GeneratedImage.get_version()),
                 transformer=mlx.utils.tree_unflatten(list(lora_layers.items()))['transformer'],
             )  # fmt:off
 
diff --git a/src/mflux/weights/weight_handler.py b/src/mflux/weights/weight_handler.py
index 4b162f0..830b5e5 100644
--- a/src/mflux/weights/weight_handler.py
+++ b/src/mflux/weights/weight_handler.py
@@ -14,7 +14,6 @@ class MetaData:
     quantization_level: int | None = None
     scale: float | None = None
     is_lora: bool = False
-    is_mflux: bool = False
     mflux_version: str | None = None
 
 
diff --git a/src/mflux/weights/weight_handler_lora.py b/src/mflux/weights/weight_handler_lora.py
index 22edb7d..da514f6 100644
--- a/src/mflux/weights/weight_handler_lora.py
+++ b/src/mflux/weights/weight_handler_lora.py
@@ -38,7 +38,7 @@ class WeightHandlerLoRA:
                         quantization_level=None,
                         scale=lora_scale,
                         is_lora=True,
-                        is_mflux=True if mflux_version is not None else False,
+                        mflux_version=mflux_version,
                     ),
                 )
                 lora_weights.append(weights)