diff --git a/hercules_opt/src/inline.rs b/hercules_opt/src/inline.rs
index f01b2366e8da12be265b65ce29f6cb977ea7c618..6e308274c65e0fbf14bb938e297a188f8ebaf0f3 100644
--- a/hercules_opt/src/inline.rs
+++ b/hercules_opt/src/inline.rs
@@ -210,7 +210,11 @@ fn inline_func(
             }
 
             // Finally, delete the call node.
-            edit = edit.replace_all_uses(id, old_id_to_new_id(called_return_data))?;
+            if let Node::Parameter { index } = called_func.nodes[called_return_data.idx()] {
+                edit = edit.replace_all_uses(id, args[index])?;
+            } else {
+                edit = edit.replace_all_uses(id, old_id_to_new_id(called_return_data))?;
+            }
             edit = edit.delete_node(control)?;
             edit = edit.delete_node(id)?;
 
diff --git a/juno_samples/fork_join_tests/src/gpu.sch b/juno_samples/fork_join_tests/src/gpu.sch
index 91bd6c79103eeb03468ea39a38995a9222836a20..364673cd2aedf323ca7727dc1ea8515188373a8d 100644
--- a/juno_samples/fork_join_tests/src/gpu.sch
+++ b/juno_samples/fork_join_tests/src/gpu.sch
@@ -52,6 +52,7 @@ slf(auto.test2);
 infer-schedules(auto.test2);
 fork-interchange[0, 1](auto.test2);
 
+inline(test6);
 fork-tile[32, 0, false, true](test6@loop);
 let out = fork-split(test6@loop);
 let out = auto-outline(test6);