diff --git a/cashocs/io/managers.py b/cashocs/io/managers.py
index 496a2165..2f0785f9 100644
--- a/cashocs/io/managers.py
+++ b/cashocs/io/managers.py
@@ -585,8 +585,16 @@ def _write_xdmf_step(
         mesh = function.function_space().mesh()
         comm = mesh.mpi_comm()
 
-        if function.function_space().ufl_element().family() == "Real":
-            space = fenics.FunctionSpace(mesh, "CG", 1)
+        if function.function_space().ufl_element().family() in [
+            "Real",
+            "NodalEnrichedElement",
+        ]:
+            if len(function.ufl_shape) > 0:
+                space = fenics.VectorFunctionSpace(
+                    mesh, "CG", 1, dim=function.ufl_shape[0]
+                )
+            else:
+                space = fenics.FunctionSpace(mesh, "CG", 1)
             function = fenics.interpolate(function, space)
 
         function.rename(function_name, function_name)