diff --git a/a3fe/_version.py b/a3fe/_version.py
index 260c070..f9aa3e1 100644
--- a/a3fe/_version.py
+++ b/a3fe/_version.py
@@ -1 +1 @@
-__version__ = "0.3.1"
+__version__ = "0.3.2"
diff --git a/a3fe/run/leg.py b/a3fe/run/leg.py
index 8c22fb1..c46b4cf 100644
--- a/a3fe/run/leg.py
+++ b/a3fe/run/leg.py
@@ -259,10 +259,10 @@ def setup(
                     relative_simulation_cost=self.relative_simulation_cost,
                     ensemble_size=self.ensemble_size,
                     lambda_values=cfg.lambda_values[self.leg_type][stage_type],
-                    base_dir=self.stage_input_dirs[stage_type].replace("/input", ""),
+                    base_dir=_os.path.dirname(self.stage_input_dirs[stage_type]),
                     input_dir=self.stage_input_dirs[stage_type],
-                    output_dir=self.stage_input_dirs[stage_type].replace(
-                        "input", "output"
+                    output_dir=_os.path.join(
+                        _os.path.dirname(self.stage_input_dirs[stage_type]), "output"
                     ),
                     stream_log_level=self.stream_log_level,
                 )
diff --git a/a3fe/tests/test_run.py b/a3fe/tests/test_run.py
index c2e19e2..18850e5 100644
--- a/a3fe/tests/test_run.py
+++ b/a3fe/tests/test_run.py
@@ -321,7 +321,7 @@ def _mock_run_process_inner(
     @pytest.fixture(scope="class")
     def setup_calc(mock_run_process):
         """Set up a calculation object completely from input files."""
-        with TemporaryDirectory() as dirname:
+        with TemporaryDirectory(prefix="temp_input_") as dirname:
             # Copy the example input directory to the temporary directory
             # as we'll create some new files there
             subprocess.run(
@@ -463,6 +463,17 @@ def test_setup_calc_sims(self, setup_calc):
                         base_dir_files = set(os.listdir(sim.base_dir))
                         assert base_dir_files == expected_base_files
 
+    def test_stage_output_path_replacement(self, setup_calc):
+        """Test that stage output paths are correctly derived from input paths by only
+        replacing the last 'input' with 'output'."""
+        for leg in setup_calc.legs:
+            for stage in leg.stages:
+                if leg.leg_type == a3.LegType.BOUND:
+                    expected_output_dir = os.path.join(
+                        os.path.dirname(stage.input_dir), "output"
+                    )
+                    assert stage.output_dir == expected_output_dir
+
 
 ######################## Tests Requiring SLURM ########################
 
diff --git a/docs/CHANGELOG.rst b/docs/CHANGELOG.rst
index 4ea289b..337b3dc 100644
--- a/docs/CHANGELOG.rst
+++ b/docs/CHANGELOG.rst
@@ -5,6 +5,7 @@ Change Log
 0.3.2
 ====================
 - Fix bug which caused somd.rst7 files in the ensemble equilibration directories to be incorrectly numbered in some cases.
+- Fix bug which caused the output directory to be incorrectly replaced with "output" in some cases.
 
 0.3.1
 ====================