diff --git a/crates/cubecl-macros/src/codegen_function/launch.rs b/crates/cubecl-macros/src/codegen_function/launch.rs
index 660e4f32..49ed4dcb 100644
--- a/crates/cubecl-macros/src/codegen_function/launch.rs
+++ b/crates/cubecl-macros/src/codegen_function/launch.rs
@@ -299,18 +299,6 @@ impl Codegen {
             let mut inputs: std::collections::BTreeMap<usize, std::sync::Arc<dyn core::any::Any>> = std::collections::BTreeMap::new();
             let mut outputs: std::collections::BTreeMap<usize, std::sync::Arc<dyn core::any::Any>> = std::collections::BTreeMap::new();
 
-            for mapping in self.settings.mappings.iter() {
-                if !inputs.contains_key(&mapping.pos_input) {
-                    inputs.insert(
-                        mapping.pos_input,
-                        #register_input_call(&mut builder, &self.settings, mapping.pos_input),
-                    );
-                }
-
-                let input = inputs.get(&mapping.pos_input).unwrap();
-                outputs.insert(mapping.pos_output, input.clone());
-            }
-
             #register_input
             #register_output
         };
@@ -325,6 +313,13 @@ impl Codegen {
             });
         }
 
+        tokens.extend(quote::quote! {
+            for mapping in self.settings.mappings.iter() {
+                let input = inputs.get(&mapping.pos_input).unwrap();
+                outputs.insert(mapping.pos_output, input.clone());
+            }
+        });
+
         if num_outputs > 0 {
             tokens.extend(quote::quote! {
                 for i in 0..#num_outputs {