diff --git a/docs/getting_started/getting_started.md b/docs/getting_started/getting_started.md
index bf1ec357..0e6c3e08 100644
--- a/docs/getting_started/getting_started.md
+++ b/docs/getting_started/getting_started.md
@@ -40,7 +40,6 @@ plot_results(boxes, root_dir=os.path.dirname(raster_path))
 
 We encourage users to try out a variety of patch sizes. For 0.1m data, 400-800px per window is appropriate, but it will depend on the density of tree plots. For coarser resolution tiles, >800px patch sizes have been effective, but we welcome feedback from users using a variety of spatial resolutions.
 
-
 ### Sample package data
 
 DeepForest comes with a small set of sample data that can be used to test out the provided examples. The data resides in the DeepForest data directory. Use the `get_data` helper function to locate the path to this directory, if needed.
diff --git a/tests/test_main.py b/tests/test_main.py
index 6a1da630..1d09172c 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -268,7 +268,7 @@ def test_predict_tile(m, raster_path):
 
     assert isinstance(prediction, pd.DataFrame)
     assert set(prediction.columns) == {
-        "xmin", "ymin", "xmax", "ymax", "label", "score", "image_path"
+        "xmin", "ymin", "xmax", "ymax", "label", "score", "image_path","geometry"
     }
     assert not prediction.empty
 
@@ -621,6 +621,6 @@ def test_predict_tile_with_crop_model(m, config):
     # Assert the result
     assert isinstance(result, pd.DataFrame)
     assert set(result.columns) == {
-        "xmin", "ymin", "xmax", "ymax", "label", "score", "cropmodel_label",
+        "xmin", "ymin", "xmax", "ymax", "label", "score", "cropmodel_label","geometry",
         "cropmodel_score", "image_path"
     }