diff --git a/example/diffusion/diffusion.cpp b/example/diffusion/diffusion.cpp
index e8f94651b1029ca7bba049b51cf4e032bce48699..e2ed2537ba91a918046c4e9ffeca3e3b6a836c7f 100644
--- a/example/diffusion/diffusion.cpp
+++ b/example/diffusion/diffusion.cpp
@@ -119,7 +119,7 @@ int main(int argc, char** argv) {
     if (!out.good()) throw std::runtime_error("Could not open output file for writing.");
     auto C = make_context({1, O.gpu});
     auto R = linear{O.L, O.dx, O.Xi, O.dX};
-    simulation S(R, partition_load_balance(R, C), C);
+    simulation S(R, C, partition_load_balance(R, C));
     S.add_sampler(all_probes, regular_schedule(O.ds), sampler);
     S.run(O.T, O.dt);
     out.close();
diff --git a/python/example/diffusion.py b/python/example/diffusion.py
index 6d9f06bd834049ed6b3e9212532bf815fc506d78..5198b4cb318e99a06189df9b89b91bd4716e6042 100644
--- a/python/example/diffusion.py
+++ b/python/example/diffusion.py
@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 
 import arbor as A
-import pandas as pd
 import seaborn as sns
 import matplotlib.pyplot as plt
 
diff --git a/test/unit/test_diffusion.cpp b/test/unit/test_diffusion.cpp
index 37e29222567d99f0be7d84d48d9af4e71c467dcd..4575ed7d81cd34c12029784461802721389530ee 100644
--- a/test/unit/test_diffusion.cpp
+++ b/test/unit/test_diffusion.cpp
@@ -114,7 +114,7 @@ testing::AssertionResult run(const linear& rec, const result_t exp) {
         }
     };
     auto ctx = make_context({arbenv::default_concurrency(), with_gpu});
-    auto sim = simulation{rec, partition_load_balance(rec, ctx), ctx};
+    auto sim = simulation{rec, ctx, partition_load_balance(rec, ctx)};
     sim.add_sampler(arb::all_probes, arb::regular_schedule(0.1), sampler);
     sim.run(0.11, 0.01);
     return all_near(sample_values, exp, epsilon);