From 771d69af202d87bc1e0acb80166e6308b4919118 Mon Sep 17 00:00:00 2001
From: augustin64 <me.git@augustin64.fr>
Date: Sat, 14 Jan 2023 15:28:02 +0100
Subject: [PATCH] Fix dataset memory leak

---
 src/cnn/train.c | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/src/cnn/train.c b/src/cnn/train.c
index a26220b..7ac13b1 100644
--- a/src/cnn/train.c
+++ b/src/cnn/train.c
@@ -288,11 +288,26 @@ void train(int dataset_type, char* images_file, char* labels_file, char* data_di
     }
     free(shuffle_index);
     free_network(network);
+
     #ifdef USE_MULTITHREADING
     free(tid);
     #else
     free(train_params);
     #endif
+    
+    if (dataset_type == 0) {
+        for (int i=0; i < nb_images_total; i++) {
+            for (int j=0; j < 28; j++) {
+                free(images[i][j]);
+            }
+            free(images[i]);
+        }
+        free(images);
+        free(labels);
+    } else {
+        free_dataset(dataset);
+    }
+
     end_time = omp_get_wtime();
     elapsed_time = end_time - algo_start;
     printf("\nTemps total: %0.1f s\n", elapsed_time);