Merge pull request #405 from PyAutoLabs/feature/ag-fit-quantity-dataset-model-pytree

fix: register DatasetModel pytree in _register_fit_quantity_pytrees

Author: Jammy2211

autogalaxy/quantity/model/analysis.py

@@ -78,12 +78,14 @@ def _register_fit_quantity_pytrees() -> None:
         carry the traced model arrays and ride as pytree children.
         """
         from autoarray.abstract_ndarray import register_instance_pytree
+        from autoarray.dataset.dataset_model import DatasetModel
         from autogalaxy.analysis.jax_pytrees import register_galaxies_pytree
 
         register_instance_pytree(
             FitQuantity,
             no_flatten=("dataset", "func_str", "use_mask_in_fit"),
         )
+        register_instance_pytree(DatasetModel)
         register_galaxies_pytree()
 
     def log_likelihood_function(self, instance: af.ModelInstance) -> float: