diff --git a/examples/driver.sbatch b/examples/driver.sbatch
index b87222f2..6657c7f1 100644
--- a/examples/driver.sbatch
+++ b/examples/driver.sbatch
@@ -30,12 +30,13 @@ do
     srun --cpu-bind=cores -n 1 julia --project=examples -e \
     "import CalibrateAtmos;
     atmos_config = CalibrateAtmos.get_atmos_config($i, $iteration; \"$job_id\"); 
-    CalibrateAtmos.run_forward_model(atmos_config);
-    include(\"examples/$job_id/observation_map.jl\");
-    observation_map($i, $iteration, \"$job_id\")"
+    CalibrateAtmos.run_forward_model(atmos_config)"
 done
 
 julia --project=examples -e \
-    "import CalibrateAtmos; CalibrateAtmos.update_ensemble(\"$job_id\", $iteration)"
+    "import CalibrateAtmos; 
+    include(\"examples/$job_id/observation_map.jl\");
+    observation_map($iteration, \"$job_id\");
+    CalibrateAtmos.update_ensemble(\"$job_id\", $iteration)"
 done
 
diff --git a/examples/sphere_held_suarez_rhoe_equilmoist/observation_map.jl b/examples/sphere_held_suarez_rhoe_equilmoist/observation_map.jl
index 3cf09a42..247149ed 100644
--- a/examples/sphere_held_suarez_rhoe_equilmoist/observation_map.jl
+++ b/examples/sphere_held_suarez_rhoe_equilmoist/observation_map.jl
@@ -2,6 +2,7 @@ using NetCDF
 import Statistics: mean
 import YAML
 import EnsembleKalmanProcesses: TOMLInterface
+import JLD2
 
 export observation_map
 
@@ -35,9 +36,11 @@ function time_avg(arr)
     return dropdims(mean(arr; dims); dims)
 end
 
-function observation_map(ensemble_size, iteration, job_id)
+function observation_map(iteration, job_id)
     config = YAML.load_file("examples/$job_id/ekp_config.yml")
     output_dir = config["output_dir"]
+    # ensemble_size = config["ensemble_size"]
+    ensemble_size = 1
     filename = "ta_60.0d_average.nc"
 
     dims = (19 - 2)
@@ -60,5 +63,5 @@ function observation_map(ensemble_size, iteration, job_id)
         "/",
     )
     iter_path = joinpath(member_path[1:(length(member_path) - 1)])
-    JLD2.save_object(joinpath(iter_path, "observation_map.nc"), G_ens)
+    JLD2.save_object(joinpath(iter_path, "observation_map.jld2"), G_ens)
 end
diff --git a/test/Project.toml b/test/Project.toml
index 71a1e9b9..84a24c9e 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -4,5 +4,6 @@ CalibrateAtmos = "4347a170-ebd6-470c-89d3-5c705c0cacc2"
 ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 EnsembleKalmanProcesses = "aa8a2aa5-91d8-4396-bcef-d4f2ec43552d"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"