diff --git a/peak_finder/driver.py b/peak_finder/driver.py
index 04e5758..eb6239a 100644
--- a/peak_finder/driver.py
+++ b/peak_finder/driver.py
@@ -206,6 +206,11 @@ def run(self):  # pylint: disable=R0912, R0914, too-many-statements
                 if False in masking_array:
                     survey.remove_vertices(~masking_array)
 
+                new_line_id = survey.get_entity(self.params.line_field.uid)[0]
+
+                if isinstance(new_line_id, ReferencedData):
+                    self.params.line_field = new_line_id
+
             line_field_obj = self.params.get_line_field(survey)
 
             if (
diff --git a/tests/utils_test.py b/tests/utils_test.py
index 8acbb16..7e84fb2 100644
--- a/tests/utils_test.py
+++ b/tests/utils_test.py
@@ -18,7 +18,7 @@ def test_get_ordered_survey_lines(tmp_path: Path):
     # Create line field with non-adjacent labels
     h5file_path = tmp_path / r"testOrderingSurveyLines.geoh5"
     # Create temp workspace
-    temp_ws = Workspace(h5file_path)
+    temp_ws = Workspace.create(h5file_path)
 
     x = np.c_[
         np.zeros(30), np.tile(np.arange(0, 10), 3), np.repeat(np.arange(0, 30, 10), 10)