diff --git a/src/main.cu b/src/main.cu
index dfc5f9c421ecf0bf23b47d5818845d8060886eea..8cbb1c4629035cce683dda0471ab7cffe20858e6 100644
--- a/src/main.cu
+++ b/src/main.cu
@@ -17,6 +17,7 @@ int main(int argc, char *argv[])
     }
 
     int node_count = atoi(argv[1]);
+    cudaDeviceSetLimit(cudaLimitMallocHeapSize, node_count * node_count * 2.5f * sizeof(long int));
     long int *adj_matrix = create_adj_matrix(node_count, node_count);
 
     long int *sub_dist = calculate_sub_matrix(adj_matrix, node_count);
@@ -31,13 +32,14 @@ int main(int argc, char *argv[])
         print_matrix(sub_dist, node_count, node_count);
     }
 
-    free(sub_dist);
-    free(adj_matrix);
+    cudaFree(sub_dist);
+    cudaFree(adj_matrix);
 }
 
 long int *calculate_sub_matrix(long int *matrix, int node_count)
 {
     long int *sub_dist = (long int *)malloc(node_count * node_count * sizeof(long int));
+    cudaMallocManaged(&sub_dist, node_count * node_count * sizeof(long int));
 
     for (int i = 0; i < node_count; i++)
     {
@@ -47,6 +49,8 @@ long int *calculate_sub_matrix(long int *matrix, int node_count)
         {
             set_el(sub_dist, node_count, j, i, temp_dist[j]);
         }
+
+        cudaFree(temp_dist);
     }
 
     return sub_dist;
diff --git a/src/utils/dijkstra.cu b/src/utils/dijkstra.cu
index 2949cdce79526fa14785aa56f3dad99631ec5fd9..930dd5d59c5c88174c7b3f2901a696ef1211268e 100644
--- a/src/utils/dijkstra.cu
+++ b/src/utils/dijkstra.cu
@@ -19,7 +19,8 @@ long int get_idx_min_dist(long int *dist, short *processed, int len)
 
 long int *dijkstra(long int *adj_matrix, int src, int size)
 {
-    long int *dist = (long int *)malloc(size * sizeof(long int));
+    long int *dist;
+    cudaMallocManaged(&dist, size * sizeof(long int));
     short processed[size];
 
     for (int i = 0; i < size; i++)
diff --git a/src/utils/matrix.cu b/src/utils/matrix.cu
index 3cb4c93f22253ab468a520aff3d9645504c661b2..f7b9bb0951e6f6735b212c66ec23cc8f58a03578 100644
--- a/src/utils/matrix.cu
+++ b/src/utils/matrix.cu
@@ -40,7 +40,8 @@ void print_matrix_to_file(long int *matrix, int width, int height, char *filenam
 
 long int *create_adj_matrix(int width, int height)
 {
-    long int *matrix = (long int *)malloc(width * height * sizeof(long int));
+    long int *matrix;
+    cudaMallocManaged(&matrix, width * height * sizeof(long int));
 
     for (int i = 0; i < height; i++)
     {