diff --git a/experiments/MnistExperiment.c b/experiments/MnistExperiment.c
index 7cf7cf1..98c7ec4 100644
--- a/experiments/MnistExperiment.c
+++ b/experiments/MnistExperiment.c
@@ -296,6 +296,8 @@ int main(void) {
             inferenceStats_t *validationResult = inferenceWithLoss(
                 model, sizeModel, batch->samples[0]->item, batch->samples[0]->label, CROSS_ENTROPY);
             validationLoss += validationResult->loss;
+            freeInferenceStats(validationResult);
+            freeBatch(batch);
         }
 
         validationLoss /= (float)testDatasetSize;
diff --git a/src/userApi/data_loader/DataLoaderApi.c b/src/userApi/data_loader/DataLoaderApi.c
index 22ecc0d..5b0b0d0 100644
--- a/src/userApi/data_loader/DataLoaderApi.c
+++ b/src/userApi/data_loader/DataLoaderApi.c
@@ -61,7 +61,7 @@ static batch_t *getBatch(dataLoader_t *dataLoader, size_t index) {
 
     size_t batchSize = dataLoader->batchSize;
     batch->size = batchSize;
-    batch->samples = *reserveMemory(batchSize * sizeof(sample_t));
+    batch->samples = *reserveMemory(batchSize * sizeof(sample_t *));
 
     size_t sampleIndex = index * batchSize;