diff --git a/src/virtual_stain_flow/trainers/AbstractTrainer.py b/src/virtual_stain_flow/trainers/AbstractTrainer.py
index 932faf3..82bf209 100644
--- a/src/virtual_stain_flow/trainers/AbstractTrainer.py
+++ b/src/virtual_stain_flow/trainers/AbstractTrainer.py
@@ -13,7 +13,7 @@
 from torch.utils.data import DataLoader
 
 from .trainer_protocol import TrainerProtocol
-from .trainer_utils import EarlyStopHelper, save_model
+from .trainer_utils import EarlyStopHelper, save_model, save_optimizer_state
 from ..metrics.AbstractMetrics import AbstractMetrics
 from ..engine.progress import Progress
 from ..datasets.data_split import default_random_split
@@ -42,6 +42,7 @@ def __init__(
         test_ratio: Optional[float] = 0.15,
         metrics: Dict[str, AbstractMetrics] = None,
         device: Optional[torch.device] = None,
+        epoch: Optional[int] = 0,
         early_termination_metric: Optional[str] = None,
         early_termination_mode: Literal['min', 'max'] = "min",
         **kwargs,
@@ -102,17 +103,19 @@ def __init__(
             **kwargs
         )
         self._init_state(
+            epoch,
             early_termination_metric, early_termination_mode, **kwargs)
 
     def _init_state(
         self, 
+        epoch,
         early_termination_metric: Optional[str] = None,
         early_termination_mode: Literal['min', 'max'] = "min",
         **kwargs
     ):
 
         # Epoch state
-        self._epoch = 0
+        self._epoch = epoch
         
         # Progress tracking for loss weight scheduling
         self._progress = Progress(epoch=0, step=0)
@@ -183,7 +186,6 @@ def _init_data(
                 **kwargs
             )
 
-            self._batch_size = batch_size
             self._train_ratio, self._val_ratio, self._test_ratio = (
                 train_ratio, val_ratio, test_ratio
             )
@@ -194,6 +196,11 @@ def _init_data(
                 "or provide at least train_loader."
             )
 
+        self._batch_size = self._train_loader.batch_size if hasattr(self._train_loader, 'batch_size') else None
+        self._train_n = len(self._train_loader.dataset) if hasattr(self._train_loader, 'dataset') else None
+        self._val_n = len(self._val_loader.dataset) if hasattr(self._val_loader, 'dataset') else None
+        self._test_n = len(self._test_loader.dataset) if hasattr(self._test_loader, 'dataset') else None
+
         return None
 
     @abstractmethod
@@ -323,7 +330,6 @@ def train(
         if hasattr(logger, "on_train_start"):
             logger.on_train_start()
 
-        self._epochs = epochs
         self._epoch_pbar: Optional[tqdm] = tqdm(
             range(epochs), desc="Training", unit="epoch") if verbose else None
         iterable = self._epoch_pbar if self._epoch_pbar else range(epochs)
@@ -420,6 +426,30 @@ def save_model(
             save_best_model=best_model
         )
 
+    def save_optimizer_state(
+        self, 
+        save_path: pathlib.Path, 
+        file_name_prefix: Optional[str] = None, 
+        file_name_suffix: Optional[str] = None, 
+        file_ext: str = '.pth',
+        recent: bool = True
+    ) -> Optional[List[pathlib.Path]]:
+        """
+        Save the optimizer state to the specified path.
+        """
+        if not recent:
+            raise NotImplementedError(
+                "Saving non-recent optimizer states is not implemented yet."
+            )
+        file_name_suffix = file_name_suffix or 'recent'
+        return save_optimizer_state(
+            trainer=self,
+            save_path=save_path,
+            file_name_prefix=file_name_prefix,
+            file_name_suffix=file_name_suffix,
+            file_ext=file_ext
+        )
+
     """
     Log property
     """
@@ -452,6 +482,18 @@ def val_ratio(self):
     @property
     def test_ratio(self):
         return self._test_ratio
+
+    @property
+    def train_n(self):
+        return self._train_n
+
+    @property
+    def val_n(self):
+        return self._val_n
+
+    @property
+    def test_n(self):
+        return self._test_n
     
     @property
     def model(self):
@@ -468,11 +510,7 @@ def device(self):
     @property
     def batch_size(self):
         return self._batch_size
-    
-    @property
-    def epochs(self):
-        return self._epochs
-    
+        
     @property
     def patience(self):
         return self._patience
diff --git a/src/virtual_stain_flow/trainers/logging_gan_trainer.py b/src/virtual_stain_flow/trainers/logging_gan_trainer.py
index e51cc14..ca9a25e 100644
--- a/src/virtual_stain_flow/trainers/logging_gan_trainer.py
+++ b/src/virtual_stain_flow/trainers/logging_gan_trainer.py
@@ -84,7 +84,6 @@ def __init__(
         super().__init__(
             model=generator, # register generator as main model for early stopping
             optimizer=generator_optimizer,
-            losses=generator_loss_group,
             device=device,
             **kwargs
         )
diff --git a/src/virtual_stain_flow/trainers/trainer_protocol.py b/src/virtual_stain_flow/trainers/trainer_protocol.py
index 5fc3cb0..e0de924 100644
--- a/src/virtual_stain_flow/trainers/trainer_protocol.py
+++ b/src/virtual_stain_flow/trainers/trainer_protocol.py
@@ -44,6 +44,18 @@ def train(self, *args: Any, **kwargs: Any) -> None: ...
     @property
     def epoch(self) -> int: ...
 
+    @property
+    def batch_size(self) -> int: ...
+
+    @property
+    def train_n(self) -> int: ...
+
+    @property
+    def val_n(self) -> int: ...
+
+    @property
+    def test_n(self) -> int: ...
+
     @property
     def device(self) -> torch.device: ...
 
@@ -59,9 +71,19 @@ def best_model(self) -> Optional[torch.nn.Module]: ...
     def save_model(
         self, 
         save_path: pathlib.Path, 
-        file_name_prefix: Optional[str], 
-        file_name_suffix: Optional[str], 
+        file_name_prefix: Optional[str] = None, 
+        file_name_suffix: Optional[str] = None, 
         file_ext: str = '.pth',
         best_model: bool = True,
     ) -> Optional[List[pathlib.Path]]:
         ...
+
+    def save_optimizer_state(
+        self, 
+        save_path: pathlib.Path, 
+        file_name_prefix: Optional[str] = None, 
+        file_name_suffix: Optional[str] = None, 
+        file_ext: str = '.pth',
+        recent: bool = True,
+    ) -> Optional[List[pathlib.Path]]:
+        ...
diff --git a/src/virtual_stain_flow/trainers/trainer_utils/__init__.py b/src/virtual_stain_flow/trainers/trainer_utils/__init__.py
index f0a09fa..a706015 100644
--- a/src/virtual_stain_flow/trainers/trainer_utils/__init__.py
+++ b/src/virtual_stain_flow/trainers/trainer_utils/__init__.py
@@ -7,10 +7,12 @@
     _get_latest_metric_value,
 )
 from .save_model import save_model
+from .save_optimizer import save_optimizer_state
 
 
 __all__ = [
     "EarlyStopHelper",
     "_get_latest_metric_value",
     "save_model",
+    "save_optimizer_state",
 ]
diff --git a/src/virtual_stain_flow/trainers/trainer_utils/save_model.py b/src/virtual_stain_flow/trainers/trainer_utils/save_model.py
index 528d81b..8063199 100644
--- a/src/virtual_stain_flow/trainers/trainer_utils/save_model.py
+++ b/src/virtual_stain_flow/trainers/trainer_utils/save_model.py
@@ -6,12 +6,15 @@
 def save_model(
     trainer: 'TrainerProtocol',
     save_path: Path,
-    file_name_prefix: str = 'generator',
+    file_name_prefix: Optional[str] = None,
     file_name_suffix: Optional[str] = None,
     file_ext: str = '.pth',
     save_best_model: bool = True,
 ) -> List[Path]:
 
+    if file_name_prefix is None:
+        file_name_prefix = 'generator'
+
     if file_name_suffix is None:
         file_name_suffix = 'weights_' + (
             'best' if save_best_model else str(trainer.epoch)
diff --git a/src/virtual_stain_flow/trainers/trainer_utils/save_optimizer.py b/src/virtual_stain_flow/trainers/trainer_utils/save_optimizer.py
new file mode 100644
index 0000000..97e899c
--- /dev/null
+++ b/src/virtual_stain_flow/trainers/trainer_utils/save_optimizer.py
@@ -0,0 +1,39 @@
+from pathlib import Path
+from typing import Optional, List
+
+import torch
+
+from ..trainer_protocol import TrainerProtocol
+
+
+def save_optimizer_state(
+    trainer: 'TrainerProtocol',
+    save_path: Path,
+    file_name_prefix: Optional[str] = None,
+    file_name_suffix: Optional[str] = None,
+    file_ext: str = '.pth',
+) -> List[Path]:
+
+    if file_name_prefix is None:
+        file_name_prefix = 'optimizer'
+
+    if file_name_suffix is None:
+        file_name_suffix = f"{trainer.epoch}"
+
+    optimizer = trainer.optimizer
+
+    if optimizer is None:
+        return []
+
+    save_file = save_path / f"{file_name_prefix}_{file_name_suffix}{file_ext}"
+
+    torch.save(
+        optimizer.state_dict(), 
+        save_file
+    )
+
+    if save_file.exists():
+        return [save_file]
+
+    return []
+ 
diff --git a/src/virtual_stain_flow/vsf_logging/MlflowLogger.py b/src/virtual_stain_flow/vsf_logging/MlflowLogger.py
index 8933283..c534378 100644
--- a/src/virtual_stain_flow/vsf_logging/MlflowLogger.py
+++ b/src/virtual_stain_flow/vsf_logging/MlflowLogger.py
@@ -15,7 +15,9 @@
     AutoLossGroupConfigLogger,
     AutoModelConfigLogger,
     AutoOptimizerConfigLogger,
+    AutoTrainerLogger
 )
+from .logger_utils import _log_artifact, _log_trainer_artifact
 from .callbacks.LoggerCallback import (
     AbstractLoggerCallback,
     log_type
@@ -145,6 +147,7 @@ def __init__(
         self._model_config_logger = AutoModelConfigLogger(self)
         self._optimizer_config_logger = AutoOptimizerConfigLogger(self)
         self._loss_group_config_logger = AutoLossGroupConfigLogger(self)
+        self._trainer_logger = AutoTrainerLogger(self)
 
         return None
     
@@ -208,7 +211,7 @@ def on_train_start(self):
         self._model_config_logger.log_model_configs(self.trainer)
         self._optimizer_config_logger.log_optimizer_configs(self.trainer)
         self._loss_group_config_logger.log_loss_group_configs(self.trainer)
-
+        self._trainer_logger.log_trainer_config(self.trainer)
         
         for callback in self.callbacks:
             # TODO consider if we want hasattr checks
@@ -248,16 +251,10 @@ def on_epoch_end(self):
 
         if self._save_model_every_n_epochs is not None:
             if self.trainer.epoch % self._save_model_every_n_epochs == 0:
-                self._save_model_weights(
-                    artifact_path='weights',
-                    best_model=False
-                )
+                _log_trainer_artifact(self.trainer, best_model=False)
 
         if self._save_best_model:
-            self._save_model_weights(
-                artifact_path='weights',
-                best_model=True
-            )
+            _log_trainer_artifact(self.trainer, best_model=True)
 
         # Call on_epoch_end for all registered callbacks
         for callback in self.callbacks:
@@ -323,14 +320,14 @@ def end_run(self):
             print("No active MLflow run to end.")
 
     """
-    Exposed? logging methods
+    Exposed logging methods
     """    
     def log_artifact(
-            self,
-            tag: str,
-            file_path: pathlib.Path,
-            stage: Optional[str] = None
-        ):
+        self,
+        tag: str,
+        file_path: pathlib.Path,
+        stage: Optional[str] = None
+    ):
         """
         Log an artifact to MLflow.
 
@@ -339,29 +336,11 @@ def log_artifact(
         :param stage: Optional stage to categorize the artifact, defaults to None.
         :raises TypeError: If file_path is not a pathlib.Path instance.
         """
-
-        if not isinstance(file_path, pathlib.Path):
-            raise TypeError("file_path must be a pathlib.Path instance.")
-
-        artifact_path = ''
-        artifact_ext = file_path.suffix.lower()
-        if artifact_ext in ['.png', '.jpg', '.jpeg', '.pdf', '.svg']:
-            # log as plot artifact
-            artifact_path += 'plots/'
-        elif artifact_ext in ['.pth', '.pt']:
-            # log as model artifact
-            artifact_path += 'weights/'
-        else:
-            # log as generic artifact
-            artifact_path += 'artifacts/'
-        
-        if stage is not None:
-            artifact_path += f"{stage}/"
-        artifact_path += f"{tag}"
-
-        mlflow.log_artifact(
-            str(file_path), 
-            artifact_path=artifact_path
+        log_subdirs = [stage] if stage is not None else []
+        log_subdirs = log_subdirs + [tag] if tag is not None else log_subdirs
+        _log_artifact(
+            file_path=file_path,
+            artifact_subdirs=log_subdirs
         )
 
     def log_metric(
@@ -464,31 +443,6 @@ def _log_callback_output(
                 continue
                 # raise TypeError("Unsupported callback return type for logging.")
 
-    def _save_model_weights(
-        self,
-        prefix: Optional[str] = None,
-        suffix: Optional[str] = None,
-        artifact_path: str = "weights",
-        best_model: bool = True
-    ):
-        with tempfile.TemporaryDirectory() as tmpdirname:
-            
-            tmpdirpath = pathlib.Path(tmpdirname)
-            
-            saved_file_paths = self.trainer.save_model(
-                save_path=tmpdirpath,
-                file_name_prefix=prefix,
-                file_name_suffix=suffix,
-                file_ext='.pth',
-                best_model=best_model
-            )
-
-            for saved_file_path in (saved_file_paths or []):
-                mlflow.log_artifact(
-                    str(saved_file_path), 
-                    artifact_path=artifact_path
-                )
-
     def log_config(
         self,
         tag: str,
@@ -550,7 +504,6 @@ def __check_trainer_bound(
         if self.trainer is None:
             raise RuntimeError("No trainer bound to logger. Cannot access trainer attributes.")
 
-
     def get_epoch(
             self
         ) -> int:
@@ -595,48 +548,8 @@ def get_model(
     @property
     def run_id(self):
         return self._run_id
-    
-    """
-    Unimplemented helper that might be useful
-    """
-    def _log_dict_as_yaml(
-        self,
-        dict: Dict[str, Any],
-    ):
-        """
-        Log a dictionary as a YAML file in MLflow.
-        
-        :param dict: The dictionary to log.
-        """
-        
-        # TODO implement this method to convert dict to YAML and log it
-        raise NotImplementedError(
-            "log_dict_as_yaml method is not implemented. "
-            "Please implement this method to log dictionary as YAML."
-        )
-    
-    def _log_dict_as_param(
-        self
-    ):
-        """
-        Log a dictionary as parameters in MLflow.
-        
-        :return: None
-        :raises NotImplementedError: If the method is not implemented.
-        
-        This method is intended to log a dictionary as parameters in MLflow.
-        It is currently not implemented and raises a NotImplementedError.
-        """
-
-        # 1. DO flatten dict
-
-        # 2. DO mlflow log
-        
-        raise NotImplementedError(
-            "log_dict_as_param method is not implemented. "
-            "Please implement this method to log dictionary as parameters."
-        )
 
+    
     """
     Overridden destructor method to ensure MLflow run is ended
     """
diff --git a/src/virtual_stain_flow/vsf_logging/auto_loggers/__init__.py b/src/virtual_stain_flow/vsf_logging/auto_loggers/__init__.py
index 2860937..2625d41 100644
--- a/src/virtual_stain_flow/vsf_logging/auto_loggers/__init__.py
+++ b/src/virtual_stain_flow/vsf_logging/auto_loggers/__init__.py
@@ -1,9 +1,11 @@
 from .loss_group_config_logger import AutoLossGroupConfigLogger
 from .model_config_logger import AutoModelConfigLogger
 from .optimizer_config_logger import AutoOptimizerConfigLogger
+from .trainer_config_logger import AutoTrainerLogger
 
 __all__ = [
     "AutoModelConfigLogger",
     "AutoOptimizerConfigLogger",
     "AutoLossGroupConfigLogger",
+    "AutoTrainerLogger",
 ]
diff --git a/src/virtual_stain_flow/vsf_logging/auto_loggers/optimizer_config_logger.py b/src/virtual_stain_flow/vsf_logging/auto_loggers/optimizer_config_logger.py
index fa3f344..c95537e 100644
--- a/src/virtual_stain_flow/vsf_logging/auto_loggers/optimizer_config_logger.py
+++ b/src/virtual_stain_flow/vsf_logging/auto_loggers/optimizer_config_logger.py
@@ -1,4 +1,5 @@
 from typing import Any, Dict, List, Optional
+import inspect
 
 import mlflow
 from torch.optim import Optimizer
@@ -44,13 +45,22 @@ def log_optimizer_configs(
                 continue
 
             try:
+                
+                defaults = dict(optimizer.defaults)
+                init_signature = inspect.signature(optimizer.__class__.__init__)
+                valid_params = init_signature.parameters.keys()
+                
                 opt_config: Optional[Dict[str, Any]] = {
                     "class_path": (
                         f"{optimizer.__class__.__module__}."
                         f"{optimizer.__class__.__name__}"
                     ),
-                    "defaults": dict(optimizer.defaults),
+                    "defaults": defaults,
+                    "init": {
+                        k: v for k, v in defaults.items() if k in valid_params
+                    }
                 }
+
             except Exception as e:
                 print(f"Could not get optimizer config for logging: {e}")
                 opt_config = None
diff --git a/src/virtual_stain_flow/vsf_logging/auto_loggers/trainer_config_logger.py b/src/virtual_stain_flow/vsf_logging/auto_loggers/trainer_config_logger.py
new file mode 100644
index 0000000..cf6c292
--- /dev/null
+++ b/src/virtual_stain_flow/vsf_logging/auto_loggers/trainer_config_logger.py
@@ -0,0 +1,35 @@
+from typing import Any, Optional
+
+from ...trainers.trainer_protocol import TrainerProtocol
+
+
+class AutoTrainerLogger:
+    """
+    Auto-log trainer metadata to MLflow.
+    """
+
+    def __init__(self, logger: Any) -> None:
+        self._logger = logger
+
+    def log_trainer_config(self, trainer: Optional[TrainerProtocol]) -> None:
+        
+        if trainer is None:
+            return
+
+        config = {
+            "class_path": f"{trainer.__class__.__module__}.{trainer.__class__.__name__}",
+            "device": str(trainer.device), # device used for training
+            "batch_size": trainer.batch_size, # batch size used for training
+            "train_n": trainer.train_n,
+            "val_n": trainer.val_n,
+            "test_n": trainer.test_n,
+        }
+
+        try:
+            self._logger.log_config(
+                tag="trainer",
+                config=config,
+                stage=None,
+            )
+        except Exception as e:
+            print(f"Could not log trainer config: {e}")
diff --git a/src/virtual_stain_flow/vsf_logging/logger_utils/__init__.py b/src/virtual_stain_flow/vsf_logging/logger_utils/__init__.py
new file mode 100644
index 0000000..04ad9c8
--- /dev/null
+++ b/src/virtual_stain_flow/vsf_logging/logger_utils/__init__.py
@@ -0,0 +1,6 @@
+from .log_artifacts import _log_artifact, _log_trainer_artifact
+
+__all__ = [
+    "_log_artifact",
+    "_log_trainer_artifact",
+]
diff --git a/src/virtual_stain_flow/vsf_logging/logger_utils/log_artifacts.py b/src/virtual_stain_flow/vsf_logging/logger_utils/log_artifacts.py
new file mode 100644
index 0000000..199fe7c
--- /dev/null
+++ b/src/virtual_stain_flow/vsf_logging/logger_utils/log_artifacts.py
@@ -0,0 +1,83 @@
+from typing import Optional, List
+import pathlib
+import tempfile
+
+import mlflow
+
+from virtual_stain_flow.trainers.trainer_protocol import TrainerProtocol
+
+
+def _log_artifact(
+    file_path: pathlib.Path,
+    artifact_path: Optional[str] = None,
+    artifact_subdirs: Optional[List[str]] = None
+) -> None:
+    """
+    Logs a single artifact to MLflow.
+
+    :param file_path: The path to the file to log as an artifact.
+    :param artifact_path: Optional artifact path within the MLflow run, defaults to None.
+    :param artifact_subdirs: Optional list of subdirectories to include in the artifact path, defaults to None.
+    :raises TypeError: If file_path is not a pathlib.Path instance.
+    """
+
+    if not isinstance(file_path, pathlib.Path):
+        raise TypeError("file_path must be a pathlib.Path instance.")
+
+    if artifact_path is None:
+        artifact_ext = file_path.suffix.lower()
+        if artifact_ext in ['.png', '.jpg', '.jpeg', '.pdf', '.svg']:
+            # log as plot artifact
+            artifact_path = 'plots'
+        elif artifact_ext in ['.pth', '.pt']:
+            # log as model artifact
+            artifact_path = 'weights'
+        else:
+            # log as generic artifact
+            artifact_path = 'artifacts'
+
+    artifact_subdirs = [] if artifact_subdirs is None else artifact_subdirs
+
+    path_parts = [artifact_path, *artifact_subdirs]
+    clean_parts = [
+        part.replace('\\', '/').strip('/')
+        for part in path_parts
+        if part and part.strip('/\\')
+    ]
+
+    # Build a lexical POSIX path for MLflow without resolving it on the host OS.
+    artifact_path = str(pathlib.PurePosixPath(*clean_parts)) if clean_parts else ''
+
+    mlflow.log_artifact(str(file_path), artifact_path=artifact_path)
+
+
+def _log_trainer_artifact(
+    trainer: "TrainerProtocol",
+    best_model: bool = True,
+) -> None:
+    """
+    Logs the model and optimizer state artifacts from a trainer to MLflow.
+    The most recent optimizer state is saved and logged regardless of the best_model flag.
+
+    :param trainer: The trainer instance adhering to TrainerProtocol.
+    :param best_model: Whether to log only the best model, defaults to True.
+    :raises TypeError: If the provided trainer does not adhere to TrainerProtocol.
+    """
+
+    if not isinstance(trainer, TrainerProtocol):
+        raise TypeError("The provided trainer must adhere to the TrainerProtocol.")
+    
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        
+        tmpdirpath = pathlib.Path(tmpdirname)
+        saved_model_paths = trainer.save_model(
+            save_path=tmpdirpath, best_model=best_model
+        )
+        for saved_model_path in (saved_model_paths or []):
+            _log_artifact(saved_model_path, artifact_path='weights')
+
+        saved_optimizer_paths = trainer.save_optimizer_state(
+            save_path=tmpdirpath, recent=True
+        )
+        for saved_optimizer_path in (saved_optimizer_paths or []):
+            _log_artifact(saved_optimizer_path, artifact_path='optimizer')
diff --git a/tests/trainers/test_abstract_trainer.py b/tests/trainers/test_abstract_trainer.py
index 2cb9099..0d22901 100644
--- a/tests/trainers/test_abstract_trainer.py
+++ b/tests/trainers/test_abstract_trainer.py
@@ -697,8 +697,9 @@ def test_batch_size_property_with_loader_init(self, minimal_model, minimal_optim
             device=torch.device('cpu')
         )
         
-        # When providing loaders, batch_size is set to None
-        assert trainer.batch_size is None
+        # When providing loaders, batch_size is inferred from the train loader
+        # should therefore match
+        assert trainer.batch_size is train_dataloader.batch_size
     
     def test_batch_size_property_default_value(self, minimal_model, minimal_optimizer, dataset_for_splitting):
         """Verify that batch_size property uses default value when not specified."""
diff --git a/tests/trainers/trainer_utils/test_save_model.py b/tests/trainers/trainer_utils/test_save_model.py
new file mode 100644
index 0000000..be65008
--- /dev/null
+++ b/tests/trainers/trainer_utils/test_save_model.py
@@ -0,0 +1,23 @@
+"""Standalone tests for save_model helper."""
+
+from types import SimpleNamespace
+
+from virtual_stain_flow.trainers.trainer_utils.save_model import save_model
+
+
+def test_save_model_returns_empty_when_target_model_is_none(tmp_path):
+    trainer = SimpleNamespace(model=None, best_model=None, epoch=2)
+
+    saved_paths = save_model(trainer=trainer, save_path=tmp_path, save_best_model=True)
+
+    assert saved_paths == []
+
+
+def test_save_model_saves_current_model_with_default_name(mock_model_with_save, tmp_path):
+    trainer = SimpleNamespace(model=mock_model_with_save, best_model=None, epoch=3)
+
+    saved_paths = save_model(trainer=trainer, save_path=tmp_path, save_best_model=False)
+
+    assert len(saved_paths) == 1
+    assert saved_paths[0].exists()
+    assert saved_paths[0].name == "generator_weights_3.pth"
diff --git a/tests/trainers/trainer_utils/test_save_optimizer.py b/tests/trainers/trainer_utils/test_save_optimizer.py
new file mode 100644
index 0000000..dac7d52
--- /dev/null
+++ b/tests/trainers/trainer_utils/test_save_optimizer.py
@@ -0,0 +1,23 @@
+"""Standalone tests for save_optimizer_state helper."""
+
+from types import SimpleNamespace
+
+from virtual_stain_flow.trainers.trainer_utils.save_optimizer import save_optimizer_state
+
+
+def test_save_optimizer_state_returns_empty_for_missing_optimizer(tmp_path):
+    trainer = SimpleNamespace(optimizer=None, epoch=7)
+
+    saved_paths = save_optimizer_state(trainer=trainer, save_path=tmp_path)
+
+    assert saved_paths == []
+
+
+def test_save_optimizer_state_saves_with_default_name(minimal_optimizer, tmp_path):
+    trainer = SimpleNamespace(optimizer=minimal_optimizer, epoch=7)
+
+    saved_paths = save_optimizer_state(trainer=trainer, save_path=tmp_path)
+
+    assert len(saved_paths) == 1
+    assert saved_paths[0].exists()
+    assert saved_paths[0].name == "optimizer_7.pth"