diff --git a/graph_framework/cuda_context.hpp b/graph_framework/cuda_context.hpp
index 8dfa9f438ca7edbe5b1d0fd5f2b8b063ae0d14c5..4ed324a646a6bc98c8d1344733643e860c3d6e26 100644
--- a/graph_framework/cuda_context.hpp
+++ b/graph_framework/cuda_context.hpp
@@ -358,7 +358,7 @@ namespace gpu {
                     needed_buffers.insert(output.get());
                 }
                 if (!needed_buffers.contains(output.get())) {
-                    buffers.push_back(reinterpret_cast<void *> (&kernel_arguments[input.get()]));
+                    buffers.push_back(reinterpret_cast<void *> (&kernel_arguments[output.get()]));
                     needed_buffers.insert(output.get());
                 }
             }