diff --git a/garlic/exp/fwi/common.nix b/garlic/exp/fwi/common.nix
index b220df2..cbc49a4 100644
--- a/garlic/exp/fwi/common.nix
+++ b/garlic/exp/fwi/common.nix
@@ -40,41 +40,85 @@ rec {
     ntasksPerNode = hw.cpusPerNode;
   };
 
+  srun = {nextStage, conf, ...}:
+  let
+    fwiParams = bsc.apps.fwi.params.override {
+      inherit (conf) nx ny nz;
+    };
+  in
+    stdexp.stdStages.srun {
+      inherit nextStage conf;
+      # Now we add some commands to execute before calling srun. These will
+      # only run in one rank (the first in the list of allocated nodes)
+      preSrun = ''
+        export GARLIC_FWI_SRUNDIR=$(pwd)
+        export GARLIC_FWI_EXECDIR="${conf.tempDir}/out/$GARLIC_USER/$GARLIC_UNIT/$GARLIC_RUN"
+        mkdir -p "$GARLIC_FWI_EXECDIR"
+
+        export GARLIC_FWI_PARAMS="${fwiParams}/fwi_params.txt"
+        export GARLIC_FWI_FREQ="${fwiParams}/fwi_frequencies.txt"
+
+        # We cannot change the working directory of srun, so we use a
+        # subshell to ignore the cd
+        (
+          # Generate the input dataset
+          >&2 echo "generating the input dataset"
+          cd "$GARLIC_FWI_EXECDIR"
+          ${fwiParams}/bin/ModelGenerator \
+            -m "$GARLIC_FWI_PARAMS" "$GARLIC_FWI_FREQ"
+        )
+      '';
+
+      postSrun = optionalString (conf.enableCTF) ''
+        # Save the traces
+        mv "$GARLIC_FWI_EXECDIR"/trace_* .
+      '' + ''
+        # Remove everything else
+        rm -rf "$GARLIC_FWI_EXECDIR"
+      '';
+    };
+
   exec = {nextStage, conf, ...}: stages.exec {
     inherit nextStage;
 
+    # FIXME: FWI should allow the I/O directory to be specified as a
+    # parameter
     pre = ''
-      CDIR=$(pwd)
-      EXECDIR="${conf.tempDir}/out/$GARLIC_USER/$GARLIC_UNIT/$GARLIC_RUN"
-      mkdir -p "$EXECDIR"
-      cd "$EXECDIR"
-      ln -fs ${conf.fwiInput}/InputModels InputModels || true
+      # Run fwi at the in a directory with fast local storage
+      cd "$GARLIC_FWI_EXECDIR"
+
+      echo >&2 "Current dir: $(pwd)"
+      echo >&2 "Using PARAMS=$GARLIC_FWI_PARAMS and FREQ=$GARLIC_FWI_FREQ"
     '' + optionalString (conf.enableCTF) ''
       export NANOS6_CONFIG_OVERRIDE="version.instrument=ctf"
     '';
 
     argv = [
-      "${conf.fwiInput}/fwi_params.txt"
-      "${conf.fwiInput}/fwi_frequencies.txt"
+      ''"$GARLIC_FWI_PARAMS"''
+      ''"$GARLIC_FWI_FREQ"''
     ] ++ optional (needsBlocksize conf) conf.blocksize ++ [
       "-1" # Fordward steps
       "-1" # Backward steps
       conf.ioFreq # Write/read frequency
     ];
-
-    post = ''
-      rm -rf Results || true
-    '' + optionalString (conf.enableCTF) ''
-      mv trace_* "$CDIR"
-    '';
   };
 
   apps = bsc.garlic.apps;
 
   # FWI program
-  program = {nextStage, conf, ...}: apps.fwi.solver.override {
-    inherit (conf) gitBranch fwiInput;
-  };
+  program = {nextStage, conf, ...}:
+  let
+    fwiParams = bsc.apps.fwi.params.override {
+      inherit (conf) nx ny nz;
+    };
+  in
+    apps.fwi.solver.override {
+      inherit (conf) gitBranch;
+      inherit fwiParams;
+    };
 
-  pipeline = stdexp.stdPipeline ++ [ exec program ];
+  pipeline = stdexp.stdPipelineOverride {
+    # Replace the stdandard srun stage with our own
+    overrides = { inherit srun; };
+  } ++ [ exec program ];
 }
diff --git a/garlic/exp/fwi/granularity.nix b/garlic/exp/fwi/granularity.nix
index 5b09e5a..47a95ab 100644
--- a/garlic/exp/fwi/granularity.nix
+++ b/garlic/exp/fwi/granularity.nix
@@ -37,10 +37,6 @@ let
     inherit (c) gitBranch blocksize;
     inherit (c.n) nx ny nz ntpn nodes;
 
-    fwiInput = bsc.apps.fwi.input.override {
-      inherit (c.n) nx ny nz;
-    };
-
     # Other FWI parameters
     ioFreq = -1;
 
diff --git a/garlic/exp/fwi/reuse.nix b/garlic/exp/fwi/reuse.nix
index e767281..aac1b10 100644
--- a/garlic/exp/fwi/reuse.nix
+++ b/garlic/exp/fwi/reuse.nix
@@ -60,10 +60,6 @@ let
     inherit (c) gitBranch blocksize;
     inherit (c.n) nx ny nz;
 
-    fwiInput = bsc.apps.fwi.input.override {
-      inherit (c.n) nx ny nz;
-    };
-
     # Repeat the execution of each unit several times
     loops = 10;
 
diff --git a/garlic/exp/fwi/ss.nix b/garlic/exp/fwi/ss.nix
index 2302d43..b3d5980 100644
--- a/garlic/exp/fwi/ss.nix
+++ b/garlic/exp/fwi/ss.nix
@@ -59,10 +59,6 @@ let
     inherit (c) gitBranch blocksize;
     inherit (c.n) nx ny nz;
 
-    fwiInput = bsc.apps.fwi.input.override {
-      inherit (c.n) nx ny nz;
-    };
-
     # Other FWI parameters
     ioFreq = -1;
 
diff --git a/garlic/exp/fwi/sync_io.nix b/garlic/exp/fwi/sync_io.nix
index 3498b0e..65bc0bf 100644
--- a/garlic/exp/fwi/sync_io.nix
+++ b/garlic/exp/fwi/sync_io.nix
@@ -54,10 +54,6 @@ let
     inherit (c) gitBranch blocksize;
     inherit (c.n) nx ny nz;
 
-    fwiInput = bsc.apps.fwi.input.override {
-      inherit (c.n) nx ny nz;
-    };
-
     # Other FWI parameters
     ioFreq = c.ioFreq;