diff --git a/spyro/examples/rectangle.py b/spyro/examples/rectangle.py
index 740cdd35..ac386711 100644
--- a/spyro/examples/rectangle.py
+++ b/spyro/examples/rectangle.py
@@ -2,6 +2,7 @@
 from spyro.examples.example_model import Example_model_acoustic
 from spyro.examples.example_model import Example_model_acoustic_FWI
 import firedrake as fire
+import copy
 
 rectangle_optimization_parameters = {
     "General": {
@@ -104,6 +105,11 @@
     "gradient_filename": None,
 }
 
+rectangle_dictionary_fwi = copy.deepcopy(rectangle_dictionary)
+rectangle_dictionary_fwi["inversion"] = {
+    "perform_fwi": True,  # switch to true to make a FWI
+}
+
 
 class Rectangle_acoustic(Example_model_acoustic):
     """
@@ -204,7 +210,7 @@ class Rectangle_acoustic_FWI(Example_model_acoustic_FWI):
     def __init__(
         self,
         dictionary=None,
-        example_dictionary=rectangle_dictionary,
+        example_dictionary=rectangle_dictionary_fwi,
         comm=None,
         periodic=False,
     ):
diff --git a/spyro/io/basicio.py b/spyro/io/basicio.py
index f25ed207..2ef876e5 100644
--- a/spyro/io/basicio.py
+++ b/spyro/io/basicio.py
@@ -28,7 +28,7 @@ def wrapper(*args, **kwargs):
             for snum in range(args[0].number_of_sources):
                 switch_serial_shot(args[0], snum)
                 output_list.append(func(*args, **kwargs))
-            
+
             return output_list
 
     return wrapper
@@ -305,7 +305,6 @@ def rebuild_empty_forward_solution(wave, time_steps):
         wave.forward_solution.append(fire.Function(wave.function_space))
 
 
-
 @ensemble_save_or_load
 def load_shots(Wave_obj, file_name="shots/shot_record_", shot_ids=0):
     """Load a `pickle` to a `numpy.ndarray`.
diff --git a/spyro/tools/velocity_smoother.py b/spyro/tools/velocity_smoother.py
index 44d6a057..85e91acc 100644
--- a/spyro/tools/velocity_smoother.py
+++ b/spyro/tools/velocity_smoother.py
@@ -3,9 +3,10 @@
 import segyio
 import numpy as np
 import matplotlib.pyplot as plt
+from SeismicMesh import write_velocity_model
 
 
-def smooth_velocity_field_file(input_filename, output_filename, sigma, show=False):
+def smooth_velocity_field_file(input_filename, output_filename, sigma, show=False, write_hdf5=True):
     """Smooths a velocity field using a Gaussian filter.
 
     Parameters
@@ -40,7 +41,7 @@ def smooth_velocity_field_file(input_filename, output_filename, sigma, show=Fals
 
     for i in range(ni):
         for j in range(nj):
-            if vp[i, j] < 1.51 and i < 400:
+            if vp[i, j] < 1.51:
                 vp_smooth[i, j] = vp[i, j]
 
     spec = segyio.spec()
@@ -72,4 +73,7 @@ def smooth_velocity_field_file(input_filename, output_filename, sigma, show=Fals
         ax.axis("equal")
         plt.show()
 
+    if write_hdf5:
+        write_velocity_model(output_filename, ofname=output_filename[:-5])
+
     return None
diff --git a/test/test_io.py b/test/test_io.py
index bd6c6d58..5cbf9506 100644
--- a/test/test_io.py
+++ b/test/test_io.py
@@ -31,8 +31,7 @@ def test_read_and_write_segy():
 
     vp.interpolate(c)
 
-    xi, yi, zi = spyro.io.write_function_to_grid(vp, V, 10.0 / 1000.0)
-    spyro.io.create_segy(zi, segy_file)
+    spyro.io.create_segy(vp, V, 10.0/1000.0, segy_file)
     write_velocity_model(segy_file, vp_name)
 
     model = {}
diff --git a/test/test_serialshot_fwi.py b/test/test_serialshot_fwi.py
index 963d248c..dfece0f5 100644
--- a/test/test_serialshot_fwi.py
+++ b/test/test_serialshot_fwi.py
@@ -92,7 +92,7 @@ def test_fwi(load_real_shot=False, use_rol=False):
         FWI_obj.set_real_velocity_model(conditional=cond, output=True, dg_velocity_model=False)
         FWI_obj.generate_real_shot_record(
             plot_model=True,
-            filename="True_experiment.png",
+            model_filename="True_experiment.png",
             abc_points=[(-0.5, 0.5), (-1.5, 0.5), (-1.5, 1.5), (-0.5, 1.5)]
         )
         np.save("real_shot_record", FWI_obj.real_shot_record)
diff --git a/test_polygon_fwi.py b/test_polygon_fwi.py
index 1ebbbddc..43715b19 100644
--- a/test_polygon_fwi.py
+++ b/test_polygon_fwi.py
@@ -4,6 +4,7 @@
 # debugpy.wait_for_client()
 import spyro
 import numpy as np
+import matplotlib.pyplot as plt
 
 
 def test_real_shot_record_generation_parallel():
@@ -56,5 +57,36 @@ def test_real_shot_record_generation_parallel():
     assert test
 
 
+def test_velocity_smoother_in_fwi():
+    dictionary = {}
+    dictionary["absorving_boundary_conditions"] = {
+        "pad_length": 2.0,  # True or false
+    }
+    dictionary["mesh"] = {
+        "h": 0.01,  # mesh size in km
+    }
+    dictionary["polygon_options"] = {
+        "water_layer_is_present": True,
+        "upper_layer": 2.0,
+        "middle_layer": 2.5,
+        "lower_layer": 3.0,
+        "polygon_layer_perturbation": 0.3,
+    }
+    dictionary["acquisition"] = {
+        "source_locations": spyro.create_transect((-0.1, 0.1), (-0.1, 0.9), 1),
+    }
+    fwi = spyro.examples.Polygon_acoustic_FWI(dictionary=dictionary, periodic=True)
+    spyro.io.create_segy(
+        fwi.initial_velocity_model,
+        fwi.function_space,
+        10.0/1000.0,
+        "velocity_models/true_case1.segy",
+    )
+    spyro.tools.velocity_smoother.smooth_velocity_field_file("velocity_models/true_case1.segy", "velocity_models/case1_sigma10.segy", 10, show=True, write_hdf5=True)
+    plt.savefig("velocity_models/case1_sigma10.png")
+    plt.close()
+
+
 if __name__ == "__main__":
-    test_real_shot_record_generation_parallel()
+    # test_real_shot_record_generation_parallel()
+    test_velocity_smoother_in_fwi()