diff --git a/graph_framework/cuda_context.hpp b/graph_framework/cuda_context.hpp
index 57b333ea7a401f55fd7599296ec9055137114b2f..f91090e887d439645440ae821ae9cf3ffa67bdbb 100644
--- a/graph_framework/cuda_context.hpp
+++ b/graph_framework/cuda_context.hpp
@@ -324,7 +324,7 @@ namespace gpu {
             check_error(cuModuleGetFunction(&function, module, kernel_name.c_str()), "cuModuleGetFunction");
 
             std::vector<void *> buffers;
-            std::set<graph::leaf_node<float, SAFE_MATH> *> needed_buffers;
+            std::set<graph::leaf_node<T, SAFE_MATH> *> needed_buffers;
 
             const size_t buffer_element_size = sizeof(T);
             for (auto &input : inputs) {