From e14d80b9f816f47fbdd76c17efcffae33ff759a8 Mon Sep 17 00:00:00 2001
From: Chandraveer <chandraveersingh6377@gmail.com>
Date: Wed, 25 Mar 2026 00:24:49 +0530
Subject: [PATCH 1/5] Fix grayscale channel dimension (1) lost in DataLoader
 (#306)

---
 R/dataset-mnist.R                   | 127 ++++++++++++++++------------
 R/transforms-array.R                |  11 ++-
 tests/testthat/test-dataset-mnist.R |  26 ++++--
 3 files changed, 106 insertions(+), 58 deletions(-)

diff --git a/R/dataset-mnist.R b/R/dataset-mnist.R
index 260e83c3..ce74a11b 100644
--- a/R/dataset-mnist.R
+++ b/R/dataset-mnist.R
@@ -117,7 +117,7 @@ mnist_dataset <- dataset(
       archive <- download_and_cache(r[1], prefix = class(self)[1])
       fs::file_copy(archive, destpath)
 
-      if (!tools::md5sum(destpath) == r[2])
+      if (!(tools::md5sum(destpath) == r[2]))
         runtime_error("Corrupt file! Delete the file in {archive} and try again.")
 
     }
@@ -148,8 +148,13 @@ mnist_dataset <- dataset(
   },
 
   .getitem = function(index) {
-    x <- self$data[index, ,]
-    y <- self$targets[index]
+    if (length(index) != 1) {
+      return(lapply(as.integer(index), function(i) self$.getitem(i)))
+    }
+
+    idx <- as.integer(index)
+    x <- self$data[idx, , , drop = FALSE]
+    y <- self$targets[idx]
 
     if (!is.null(self$transform))
       x <- self$transform(x)
@@ -178,7 +183,7 @@ mnist_dataset <- dataset(
 #' @describeIn mnist_dataset Kuzushiji-MNIST cursive Japanese character dataset.
 #' @export
 kmnist_dataset <- dataset(
-  name = "kminst_dataset",
+  name = "kmnist_dataset",
   inherit = mnist_dataset,
   archive_size = "21 MB",
   resources = list(
@@ -251,55 +256,63 @@ qmnist_dataset <- dataset(
   },
 
   download = function() {
-    if (self$check_exists())
-      return(NULL)
-
     fs::dir_create(self$raw_folder)
     fs::dir_create(self$processed_folder)
 
-    cli_inform("Downloading {.cls {class(self)[[1]]}} ...")
-    for (r in self$resources[[self$split]]) {
-      filename <- basename(r[1])
-      destpath <- file.path(self$raw_folder, filename)
-
-      archive <- download_and_cache(r[1], prefix = glue::glue("qmnist-{self$split}"))
-      fs::file_copy(archive, destpath, overwrite = TRUE)
-
-      if (!tools::md5sum(destpath) == r[2])
-        runtime_error("Corrupt file! Delete the file in {archive} and try again.")
+    split_exists <- function(split_name) {
+      fs::file_exists(file.path(self$processed_folder, self$files[[split_name]]))
     }
 
-    cli_inform("Processing {.cls {class(self)[[1]]}} ...")
-
-
-    if (self$split == "train") {
-      saveRDS(
-        list(
-          read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-train-images-idx3-ubyte.gz")),
-          read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-train-labels-idx2-int.gz"))
-        ),
-        file.path(self$processed_folder, self$files$train)
-      )
-    }
-
-    if (self$split == "test") {
-      saveRDS(
-        list(
-          read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-test-images-idx3-ubyte.gz")),
-          read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-test-labels-idx2-int.gz"))
-        ),
-        file.path(self$processed_folder, self$files$test)
-      )
-    }
+    if (all(vapply(names(self$resources), split_exists, logical(1))))
+      return(NULL)
 
-    if (self$split == "nist") {
-      saveRDS(
-        list(
-          read_sn3_pascalvincent(file.path(self$raw_folder, "xnist-images-idx3-ubyte.xz")),
-          read_sn3_pascalvincent(file.path(self$raw_folder, "xnist-labels-idx2-int.xz"))
-        ),
-        file.path(self$processed_folder, self$files$nist)
-      )
+    cli_inform("Downloading {.cls {class(self)[[1]]}} ...")
+    for (split_name in names(self$resources)) {
+      if (split_exists(split_name))
+        next
+
+      for (r in self$resources[[split_name]]) {
+        filename <- basename(r[1])
+        destpath <- file.path(self$raw_folder, filename)
+
+        archive <- download_and_cache(r[1], prefix = glue::glue("qmnist-{split_name}"))
+        fs::file_copy(archive, destpath, overwrite = TRUE)
+
+        if (!(tools::md5sum(destpath) == r[2]))
+          runtime_error("Corrupt file! Delete the file in {archive} and try again.")
+      }
+
+      cli_inform("Processing {.val {split_name}} split for {.cls {class(self)[[1]]}} ...")
+
+      if (split_name == "train") {
+        saveRDS(
+          list(
+            read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-train-images-idx3-ubyte.gz")),
+            read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-train-labels-idx2-int.gz"))
+          ),
+          file.path(self$processed_folder, self$files$train)
+        )
+      }
+
+      if (split_name == "test") {
+        saveRDS(
+          list(
+            read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-test-images-idx3-ubyte.gz")),
+            read_sn3_pascalvincent(file.path(self$raw_folder, "qmnist-test-labels-idx2-int.gz"))
+          ),
+          file.path(self$processed_folder, self$files$test)
+        )
+      }
+
+      if (split_name == "nist") {
+        saveRDS(
+          list(
+            read_sn3_pascalvincent(file.path(self$raw_folder, "xnist-images-idx3-ubyte.xz")),
+            read_sn3_pascalvincent(file.path(self$raw_folder, "xnist-labels-idx2-int.xz"))
+          ),
+          file.path(self$processed_folder, self$files$nist)
+        )
+      }
     }
 
     cli_inform("Dataset {.cls {class(self)[[1]]}} downloaded and extracted successfully.")
@@ -387,9 +400,11 @@ emnist_collection <- dataset(
     self$processed_folder <- file.path(root, class(self)[1], "processed")
     self$transform <- transform
     self$target_transform <- target_transform
-    self$class <- self$classes_all_dataset[[self$dataset]]
+    self$classes <- self$classes_all_dataset[[self$dataset]]
+    self$class <- self$classes
 
     if (download) {
+      cli_warn("This is a large dataset (~540 MB). Download may take a while.")
       cli_inform("{.cls {class(self)[[1]]}} (~{.emph {self$archive_size}}) will be downloaded and processed if not already available.")
       self$download()
     }
@@ -413,7 +428,7 @@ emnist_collection <- dataset(
     url <- self$resources[[1]][1]
     archive <- download_and_cache(url, prefix = class(self)[1])
 
-    if (!tools::md5sum(archive) == self$resources[[1]][2])
+    if (!(tools::md5sum(archive) == self$resources[[1]][2]))
       runtime_error("Corrupt file! Delete the file in {archive} and try again.")
 
     unzip_dir <- file.path(self$raw_folder, "unzipped")
@@ -436,9 +451,13 @@ emnist_collection <- dataset(
   },
 
   .getitem = function(index) {
+    if (length(index) != 1) {
+      return(lapply(as.integer(index), function(i) self$.getitem(i)))
+    }
 
-    x <- self$data[index, , ]
-    y <- self$targets[index]
+    idx <- as.integer(index)
+    x <- self$data[idx, , , drop = FALSE]
+    y <- self$targets[idx]
 
     if (!is.null(self$transform))
       x <- self$transform(x)
@@ -456,7 +475,11 @@ emnist_collection <- dataset(
 )
 
 read_sn3_pascalvincent <- function(path) {
-  x <- gzfile(path, open = "rb")
+  x <- if (grepl("\\\\.xz$", path, ignore.case = TRUE)) {
+    xzfile(path, open = "rb")
+  } else {
+    gzfile(path, open = "rb")
+  }
   on.exit({close(x)})
 
   magic <- readBin(x, endian = "big", what = integer(), n = 1)
diff --git a/R/transforms-array.R b/R/transforms-array.R
index 6c83d270..df19d36b 100644
--- a/R/transforms-array.R
+++ b/R/transforms-array.R
@@ -5,7 +5,16 @@ transform_to_tensor.array <- function(img) {
   if (length(dim(img)) == 2)
     dim(img) <- c(dim(img), 1)
 
-  res <- torch::torch_tensor(img)$permute(c(3, 1, 2))
+  dims <- dim(img)
+  if (length(dims) != 3)
+    stop("Expected a 2D or 3D array.")
+
+  # Support both HWC (default image arrays) and CHW (channel-first arrays).
+  if (dims[1] <= 4 && dims[3] > 4) {
+    res <- torch::torch_tensor(img)
+  } else {
+    res <- torch::torch_tensor(img)$permute(c(3, 1, 2))
+  }
 
   if (res$dtype == torch::torch_long())
     res <- res/255
diff --git a/tests/testthat/test-dataset-mnist.R b/tests/testthat/test-dataset-mnist.R
index 0819ff25..98e610a3 100644
--- a/tests/testthat/test-dataset-mnist.R
+++ b/tests/testthat/test-dataset-mnist.R
@@ -11,11 +11,21 @@ test_that("tests for the mnist dataset", {
   ds <- mnist_dataset(dir, download = TRUE)
 
   i <- ds[1]
-  expect_equal(dim(i[[1]]), c(28, 28))
+  expect_equal(dim(i[[1]]), c(1, 28, 28))
   expect_equal(i[[2]], 6)
   expect_length(ds, 60000)
 
+  raw_items <- ds$.getitem(c(1, 2))
+  expect_length(raw_items, 2)
+  expect_named(raw_items[[1]], c("x", "y"))
+  expect_equal(dim(raw_items[[1]]$x), c(1, 28, 28))
+
   ds <- mnist_dataset(dir, transform = transform_to_tensor)
+  items <- ds$.getitem(c(1, 2))
+  expect_length(items, 2)
+  expect_named(items[[1]], c("x", "y"))
+  expect_tensor_shape(items[[1]]$x, c(1, 28, 28))
+
   dl <- torch::dataloader(ds, batch_size = 32)
   expect_length(dl, 1875)
   iter <- dataloader_make_iter(dl)
@@ -25,6 +35,12 @@ test_that("tests for the mnist dataset", {
   expect_true((torch_max(i[[1]]) <= 1)$item())
   expect_named(i, c("x", "y"))
 
+  # Regression: grayscale channel must be preserved for larger batches.
+  dl_128 <- torch::dataloader(ds, batch_size = 128)
+  i_128 <- dataloader_next(dataloader_make_iter(dl_128))
+  expect_tensor_shape(i_128$x, c(128, 1, 28, 28))
+  expect_tensor_shape(i_128$y, 128)
+
 })
 
 test_that("tests for the kmnist dataset", {
@@ -36,7 +52,7 @@ test_that("tests for the kmnist dataset", {
   ds <- kmnist_dataset(dir, download = TRUE)
 
   i <- ds[1]
-  expect_equal(dim(i[[1]]), c(28, 28))
+  expect_equal(dim(i[[1]]), c(1, 28, 28))
   expect_equal(i[[2]], 6)
   expect_length(ds, 60000)
 
@@ -62,7 +78,7 @@ test_that("fashion_mnist_dataset loads correctly", {
   expect_s3_class(ds, "fashion_mnist_dataset")
   expect_type(ds$.getitem(1), "list")
   expect_named(ds$.getitem(1), c("x", "y"))
-  expect_equal(dim(as.array(ds$.getitem(1)$x)), c(28, 28))
+  expect_equal(dim(as.array(ds$.getitem(1)$x)), c(1, 28, 28))
   expect_true(ds$.getitem(1)$y >= 1 && ds$.getitem(1)$y <= 10)
 
   ds2 <- fashion_mnist_dataset(dir, transform = transform_to_tensor)
@@ -96,7 +112,7 @@ test_that("tests for the emnist dataset", {
   first_item <- emnist[1]
   expect_named(first_item, c("x", "y"))
   expect_s3_class(first_item$x, "array")
-  expect_equal(dim(first_item$x), c(28,28))
+  expect_equal(dim(first_item$x), c(1, 28, 28))
   expect_equal((first_item[[2]]), 19)
 
   emnist <- emnist_collection(dir, dataset = "bymerge", download = TRUE)
@@ -157,7 +173,7 @@ test_that("tests for the qmnist dataset", {
     ds <- qmnist_dataset(dir, split = split, download = TRUE)
 
     i <- ds[1]
-    expect_equal(dim(i[[1]]), c(28, 28))
+    expect_equal(dim(i[[1]]), c(1, 28, 28))
     expect_true(i[[2]] %in% 1:10)
 
     expect_gt(length(ds), 0)

From 1e50c94ec06f332b2ec5601478f523732be2c897 Mon Sep 17 00:00:00 2001
From: Chandraveer <chandraveersingh6377@gmail.com>
Date: Sat, 28 Mar 2026 23:40:39 +0530
Subject: [PATCH 2/5] fix requested changes

---
 NAMESPACE             |  1 +
 R/dataset-mnist.R     | 26 ++++++++++++--------------
 R/transforms-array.R  |  2 +-
 R/transforms-magick.R |  2 ++
 4 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/NAMESPACE b/NAMESPACE
index 06124d90..5c0fdcfb 100644
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -6,6 +6,7 @@ S3method(draw_bounding_boxes,torch_tensor)
 S3method(draw_segmentation_masks,default)
 S3method(draw_segmentation_masks,image_with_segmentation_mask)
 S3method(draw_segmentation_masks,torch_tensor)
+S3method(get_image_size,"magick-image")
 S3method(transform_adjust_brightness,default)
 S3method(transform_adjust_brightness,torch_tensor)
 S3method(transform_adjust_contrast,default)
diff --git a/R/dataset-mnist.R b/R/dataset-mnist.R
index 8f5ff13c..184bff9f 100644
--- a/R/dataset-mnist.R
+++ b/R/dataset-mnist.R
@@ -149,13 +149,8 @@ mnist_dataset <- dataset(
   },
 
   .getitem = function(index) {
-    if (length(index) != 1) {
-      return(lapply(as.integer(index), function(i) self$.getitem(i)))
-    }
-
-    idx <- as.integer(index)
-    x <- self$data[idx, , , drop = FALSE]
-    y <- self$targets[idx]
+    x <- self$data[index, , , drop = FALSE]
+    y <- self$targets[index]
 
     if (!is.null(self$transform))
       x <- self$transform(x)
@@ -166,6 +161,10 @@ mnist_dataset <- dataset(
     list(x = x, y = y)
   },
 
+  .getbatch = function(index) {
+    lapply(as.integer(index), function(i) self$.getitem(i))
+  },
+
   .length = function() {
     dim(self$data)[1]
   },
@@ -451,13 +450,8 @@ emnist_collection <- dataset(
   },
 
   .getitem = function(index) {
-    if (length(index) != 1) {
-      return(lapply(as.integer(index), function(i) self$.getitem(i)))
-    }
-
-    idx <- as.integer(index)
-    x <- self$data[idx, , , drop = FALSE]
-    y <- self$targets[idx]
+    x <- self$data[index, , , drop = FALSE]
+    y <- self$targets[index]
 
     if (!is.null(self$transform))
       x <- self$transform(x)
@@ -468,6 +462,10 @@ emnist_collection <- dataset(
     list(x = x, y = y)
   },
 
+  .getbatch = function(index) {
+    lapply(as.integer(index), function(i) self$.getitem(i))
+  },
+
   .length = function() {
     dim(self$data)[1]
   }
diff --git a/R/transforms-array.R b/R/transforms-array.R
index df19d36b..4c07fc3f 100644
--- a/R/transforms-array.R
+++ b/R/transforms-array.R
@@ -7,7 +7,7 @@ transform_to_tensor.array <- function(img) {
 
   dims <- dim(img)
   if (length(dims) != 3)
-    stop("Expected a 2D or 3D array.")
+    value_error("Expected a 2D or 3D array.")
 
   # Support both HWC (default image arrays) and CHW (channel-first arrays).
   if (dims[1] <= 4 && dims[3] > 4) {
diff --git a/R/transforms-magick.R b/R/transforms-magick.R
index c8fa9aff..10191618 100644
--- a/R/transforms-magick.R
+++ b/R/transforms-magick.R
@@ -53,6 +53,8 @@
 
 # Utils -------------------------------------------------------------------
 
+#' @method get_image_size magick-image
+#' @export
 `get_image_size.magick-image` <- function(img) {
   info <- magick::image_info(img)
   c(info$width, info$height)

From c6e98bde6101c99373a6797a19a3b2006c7caf11 Mon Sep 17 00:00:00 2001
From: Chandraveer <chandraveersingh6377@gmail.com>
Date: Sun, 29 Mar 2026 02:07:51 +0530
Subject: [PATCH 3/5] fix test failure

---
 R/dataset-mnist.R                   | 66 ++++++++++++++++++++++++-----
 tests/testthat/test-dataset-mnist.R | 10 +++--
 2 files changed, 62 insertions(+), 14 deletions(-)

diff --git a/R/dataset-mnist.R b/R/dataset-mnist.R
index 184bff9f..9d10ab48 100644
--- a/R/dataset-mnist.R
+++ b/R/dataset-mnist.R
@@ -149,20 +149,43 @@ mnist_dataset <- dataset(
   },
 
   .getitem = function(index) {
+    index <- as.integer(index)
     x <- self$data[index, , , drop = FALSE]
     y <- self$targets[index]
 
-    if (!is.null(self$transform))
-      x <- self$transform(x)
+    if (!is.null(self$transform)) {
+      if (length(index) > 1) {
+        x_list <- lapply(seq_along(index), function(i) {
+          self$transform(x[i, , , drop = FALSE])
+        })
+        if (length(x_list) > 0 && inherits(x_list[[1]], "torch_tensor")) {
+          x <- torch::torch_stack(x_list)
+        } else {
+          x <- x_list
+        }
+      } else {
+        x <- self$transform(x)
+      }
+    }
 
-    if (!is.null(self$target_transform))
-      y <- self$target_transform(y)
+    if (!is.null(self$target_transform)) {
+      if (length(index) > 1) {
+        y_list <- lapply(as.list(y), self$target_transform)
+        if (length(y_list) > 0 && inherits(y_list[[1]], "torch_tensor")) {
+          y <- torch::torch_stack(y_list)
+        } else {
+          y <- unlist(y_list, use.names = FALSE)
+        }
+      } else {
+        y <- self$target_transform(y)
+      }
+    }
 
     list(x = x, y = y)
   },
 
   .getbatch = function(index) {
-    lapply(as.integer(index), function(i) self$.getitem(i))
+    self$.getitem(index)
   },
 
   .length = function() {
@@ -450,20 +473,43 @@ emnist_collection <- dataset(
   },
 
   .getitem = function(index) {
+    index <- as.integer(index)
     x <- self$data[index, , , drop = FALSE]
     y <- self$targets[index]
 
-    if (!is.null(self$transform))
-      x <- self$transform(x)
+    if (!is.null(self$transform)) {
+      if (length(index) > 1) {
+        x_list <- lapply(seq_along(index), function(i) {
+          self$transform(x[i, , , drop = FALSE])
+        })
+        if (length(x_list) > 0 && inherits(x_list[[1]], "torch_tensor")) {
+          x <- torch::torch_stack(x_list)
+        } else {
+          x <- x_list
+        }
+      } else {
+        x <- self$transform(x)
+      }
+    }
 
-    if (!is.null(self$target_transform))
-      y <- self$target_transform(y)
+    if (!is.null(self$target_transform)) {
+      if (length(index) > 1) {
+        y_list <- lapply(as.list(y), self$target_transform)
+        if (length(y_list) > 0 && inherits(y_list[[1]], "torch_tensor")) {
+          y <- torch::torch_stack(y_list)
+        } else {
+          y <- unlist(y_list, use.names = FALSE)
+        }
+      } else {
+        y <- self$target_transform(y)
+      }
+    }
 
     list(x = x, y = y)
   },
 
   .getbatch = function(index) {
-    lapply(as.integer(index), function(i) self$.getitem(i))
+    self$.getitem(index)
   },
 
   .length = function() {
diff --git a/tests/testthat/test-dataset-mnist.R b/tests/testthat/test-dataset-mnist.R
index 98e610a3..decc6d2d 100644
--- a/tests/testthat/test-dataset-mnist.R
+++ b/tests/testthat/test-dataset-mnist.R
@@ -17,14 +17,16 @@ test_that("tests for the mnist dataset", {
 
   raw_items <- ds$.getitem(c(1, 2))
   expect_length(raw_items, 2)
-  expect_named(raw_items[[1]], c("x", "y"))
-  expect_equal(dim(raw_items[[1]]$x), c(1, 28, 28))
+  expect_named(raw_items, c("x", "y"))
+  expect_equal(dim(raw_items$x), c(2, 28, 28))
+  expect_equal(length(raw_items$y), 2)
 
   ds <- mnist_dataset(dir, transform = transform_to_tensor)
   items <- ds$.getitem(c(1, 2))
   expect_length(items, 2)
-  expect_named(items[[1]], c("x", "y"))
-  expect_tensor_shape(items[[1]]$x, c(1, 28, 28))
+  expect_named(items, c("x", "y"))
+  expect_tensor_shape(items$x, c(2, 1, 28, 28))
+  expect_equal(length(items$y), 2)
 
   dl <- torch::dataloader(ds, batch_size = 32)
   expect_length(dl, 1875)

From 44e4c8744a5d9641995700d78eb8a4482fc89a42 Mon Sep 17 00:00:00 2001
From: "C. Regouby" <christophe.regouby@free.fr>
Date: Sun, 29 Mar 2026 23:47:20 +0200
Subject: [PATCH 4/5] align all .getitem() to be ndim = 3, .getbatch() to be
 ndim = 4 and ds[1] to be ndim=4. make emnist_collection inherit mnist to
 remove duplicated code fix tests accordingly

---
 R/dataset-mnist.R                   | 89 +++++++----------------------
 tests/testthat/test-dataset-mnist.R | 44 ++++++++------
 2 files changed, 47 insertions(+), 86 deletions(-)

diff --git a/R/dataset-mnist.R b/R/dataset-mnist.R
index 9d10ab48..0d8e3951 100644
--- a/R/dataset-mnist.R
+++ b/R/dataset-mnist.R
@@ -149,45 +149,36 @@ mnist_dataset <- dataset(
   },
 
   .getitem = function(index) {
-    index <- as.integer(index)
     x <- self$data[index, , , drop = FALSE]
     y <- self$targets[index]
 
     if (!is.null(self$transform)) {
-      if (length(index) > 1) {
-        x_list <- lapply(seq_along(index), function(i) {
-          self$transform(x[i, , , drop = FALSE])
-        })
-        if (length(x_list) > 0 && inherits(x_list[[1]], "torch_tensor")) {
-          x <- torch::torch_stack(x_list)
-        } else {
-          x <- x_list
-        }
-      } else {
         x <- self$transform(x)
-      }
     }
 
     if (!is.null(self$target_transform)) {
-      if (length(index) > 1) {
-        y_list <- lapply(as.list(y), self$target_transform)
-        if (length(y_list) > 0 && inherits(y_list[[1]], "torch_tensor")) {
-          y <- torch::torch_stack(y_list)
-        } else {
-          y <- unlist(y_list, use.names = FALSE)
-        }
-      } else {
         y <- self$target_transform(y)
-      }
     }
 
     list(x = x, y = y)
   },
 
   .getbatch = function(index) {
-    self$.getitem(index)
-  },
+    x_flat <- self$data[index, , , drop = FALSE]
+    # unsqueeze(2) on array
+    x <- array(x_flat, dim = c(length(index), 1, dim(x_flat)[2:3]))
+    y <- self$targets[index]
+
+    if (!is.null(self$transform)) {
+      x <- torch::torch_stack(lapply(seq_len(dim(x)[1]), function(i) self$transform(x[i,,,])))
+    }
+
+    if (!is.null(self$target_transform)) {
+      y <- self$target_transform(y)
+    }
 
+    list(x = x, y = y)
+  },
   .length = function() {
     dim(self$data)[1]
   },
@@ -380,6 +371,7 @@ fashion_mnist_dataset <- dataset(
 #' @export
 emnist_collection <- dataset(
   name = "emnist_collection",
+  inherit = mnist_dataset,
   archive_size = "540 MB",
 
   resources = list(
@@ -419,8 +411,6 @@ emnist_collection <- dataset(
     self$split <- match.arg(split, choices = c("train", "test"))
     self$dataset <- match.arg(dataset,  choices = names(self$classes_all_dataset))
     self$root_path <- root
-    self$raw_folder <- file.path(root, class(self)[1], "raw")
-    self$processed_folder <- file.path(root, class(self)[1], "processed")
     self$transform <- transform
     self$target_transform <- target_transform
     self$classes <- self$classes_all_dataset[[self$dataset]]
@@ -472,50 +462,15 @@ emnist_collection <- dataset(
     fs::file_exists(file.path(self$processed_folder, self$rds_file(self$split, self$dataset)))
   },
 
-  .getitem = function(index) {
-    index <- as.integer(index)
-    x <- self$data[index, , , drop = FALSE]
-    y <- self$targets[index]
-
-    if (!is.null(self$transform)) {
-      if (length(index) > 1) {
-        x_list <- lapply(seq_along(index), function(i) {
-          self$transform(x[i, , , drop = FALSE])
-        })
-        if (length(x_list) > 0 && inherits(x_list[[1]], "torch_tensor")) {
-          x <- torch::torch_stack(x_list)
-        } else {
-          x <- x_list
-        }
-      } else {
-        x <- self$transform(x)
-      }
-    }
+  active = list(
+    raw_folder = function() {
+      file.path(self$root_path, "emnist_collection", "raw")
+    },
 
-    if (!is.null(self$target_transform)) {
-      if (length(index) > 1) {
-        y_list <- lapply(as.list(y), self$target_transform)
-        if (length(y_list) > 0 && inherits(y_list[[1]], "torch_tensor")) {
-          y <- torch::torch_stack(y_list)
-        } else {
-          y <- unlist(y_list, use.names = FALSE)
-        }
-      } else {
-        y <- self$target_transform(y)
-      }
+    processed_folder = function() {
+      file.path(self$root_path, "emnist_collection", "processed")
     }
-
-    list(x = x, y = y)
-  },
-
-  .getbatch = function(index) {
-    self$.getitem(index)
-  },
-
-  .length = function() {
-    dim(self$data)[1]
-  }
-
+  )
 )
 
 read_sn3_pascalvincent <- function(path) {
diff --git a/tests/testthat/test-dataset-mnist.R b/tests/testthat/test-dataset-mnist.R
index decc6d2d..aeb7d8a0 100644
--- a/tests/testthat/test-dataset-mnist.R
+++ b/tests/testthat/test-dataset-mnist.R
@@ -11,22 +11,28 @@ test_that("tests for the mnist dataset", {
   ds <- mnist_dataset(dir, download = TRUE)
 
   i <- ds[1]
-  expect_equal(dim(i[[1]]), c(1, 28, 28))
+  expect_equal(dim(i[[1]]), c(1, 1, 28, 28))
   expect_equal(i[[2]], 6)
   expect_length(ds, 60000)
 
-  raw_items <- ds$.getitem(c(1, 2))
-  expect_length(raw_items, 2)
-  expect_named(raw_items, c("x", "y"))
-  expect_equal(dim(raw_items$x), c(2, 28, 28))
-  expect_equal(length(raw_items$y), 2)
+  raw_item <- ds$.getitem(5)
+  expect_length(raw_item, 2)
+  expect_named(raw_item, c("x", "y"))
+  expect_equal(dim(raw_item$x), c(1, 28, 28))
+  expect_equal(length(raw_item$y), 1)
+
+  raw_batch <- ds$.getbatch(c(1, 2))
+  expect_length(raw_batch, 2)
+  expect_named(raw_batch, c("x", "y"))
+  expect_equal(dim(raw_batch$x), c(2, 1, 28, 28))
+  expect_equal(length(raw_batch$y), 2)
 
   ds <- mnist_dataset(dir, transform = transform_to_tensor)
-  items <- ds$.getitem(c(1, 2))
-  expect_length(items, 2)
-  expect_named(items, c("x", "y"))
-  expect_tensor_shape(items$x, c(2, 1, 28, 28))
-  expect_equal(length(items$y), 2)
+  batch <- ds$.getbatch(c(1, 2))
+  expect_length(batch, 2)
+  expect_named(batch, c("x", "y"))
+  expect_tensor_shape(batch$x, c(2, 1, 28, 28))
+  expect_equal(length(batch$y), 2)
 
   dl <- torch::dataloader(ds, batch_size = 32)
   expect_length(dl, 1875)
@@ -54,7 +60,7 @@ test_that("tests for the kmnist dataset", {
   ds <- kmnist_dataset(dir, download = TRUE)
 
   i <- ds[1]
-  expect_equal(dim(i[[1]]), c(1, 28, 28))
+  expect_equal(dim(i[[1]]), c(1, 1, 28, 28))
   expect_equal(i[[2]], 6)
   expect_length(ds, 60000)
 
@@ -80,7 +86,7 @@ test_that("fashion_mnist_dataset loads correctly", {
   expect_s3_class(ds, "fashion_mnist_dataset")
   expect_type(ds$.getitem(1), "list")
   expect_named(ds$.getitem(1), c("x", "y"))
-  expect_equal(dim(as.array(ds$.getitem(1)$x)), c(1, 28, 28))
+  expect_equal(dim(ds$.getitem(1)$x), c(1, 28, 28))
   expect_true(ds$.getitem(1)$y >= 1 && ds$.getitem(1)$y <= 10)
 
   ds2 <- fashion_mnist_dataset(dir, transform = transform_to_tensor)
@@ -106,14 +112,14 @@ test_that("tests for the emnist dataset", {
   expect_length(emnist, 18800)
   first_item <- emnist[1]
   expect_named(first_item, c("x", "y"))
-  expect_s3_class(first_item$x, "array")
+  expect_true(is.array(first_item$x))
   expect_equal((first_item[[2]]), 42)
 
   emnist <- emnist_collection(dir, dataset = "byclass", split = "test", download = TRUE)
   expect_length(emnist, 116323)
   first_item <- emnist[1]
   expect_named(first_item, c("x", "y"))
-  expect_s3_class(first_item$x, "array")
+  expect_true(is.array(first_item$x))
   expect_equal(dim(first_item$x), c(1, 28, 28))
   expect_equal((first_item[[2]]), 19)
 
@@ -121,7 +127,7 @@ test_that("tests for the emnist dataset", {
   expect_length(emnist, 116323)
   first_item <- emnist[1]
   expect_named(first_item, c("x", "y"))
-  expect_s3_class(first_item$x, "array")
+  expect_true(is.array(first_item$x))
   expect_equal((first_item[[2]]), 25)
 
   emnist <- emnist_collection(dir, dataset = "letters", split = "train", download = TRUE,
@@ -137,14 +143,14 @@ test_that("tests for the emnist dataset", {
   expect_length(emnist, 40000)
   first_item <- emnist[1]
   expect_named(first_item, c("x", "y"))
-  expect_s3_class(first_item$x, "array")
+  expect_true(is.array(first_item$x))
   expect_equal((first_item[[2]]), 1)
 
   emnist <- emnist_collection(dir, dataset = "mnist", split = "train", download = TRUE)
   expect_length(emnist, 60000)
   first_item <- emnist[1]
   expect_named(first_item, c("x", "y"))
-  expect_s3_class(first_item$x, "array")
+  expect_true(is.array(first_item$x))
   expect_equal((first_item[[2]]), 5)
 
   ds2 <- emnist_collection(
@@ -175,7 +181,7 @@ test_that("tests for the qmnist dataset", {
     ds <- qmnist_dataset(dir, split = split, download = TRUE)
 
     i <- ds[1]
-    expect_equal(dim(i[[1]]), c(1, 28, 28))
+    expect_equal(dim(i[[1]]), c(1, 1, 28, 28))
     expect_true(i[[2]] %in% 1:10)
 
     expect_gt(length(ds), 0)

From ec658a09273c5b0e12a66fb202861bef8dd59ebf Mon Sep 17 00:00:00 2001
From: "C. Regouby" <christophe.regouby@free.fr>
Date: Mon, 30 Mar 2026 00:29:23 +0200
Subject: [PATCH 5/5] add NEWS

---
 NEWS.md | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/NEWS.md b/NEWS.md
index 932c3c6a..83f12b5b 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -12,6 +12,7 @@
 * Added article showcasing `model_fcn_resnet50()` with visualization utilities `draw_segmentation_masks()` and `vision_make_grid()` (@DerrickUnleashed, #281).
 * Added collection dataset catalog with `search_collection()`, `get_collection_catalog()`, and `list_collection_datasets()` functions for discovering and exploring collections (#271, @ANAMASGARD).
 * Added `target_transform_coco_masks()` and `target_transform_trimap_masks()` transformation functions for explicit segmentation mask generation (@ANAMASGARD).
+* Added support for `connectivity` argument for drawing lines between keypoints in `draw_keypoints()` (@DerrickUnleashed #303)
 
 ## New models
 
@@ -26,11 +27,15 @@
 
 ## Bug fixes and improvements
 
+* `mnist_datataset()` and derivatives now correctly return item x() values with a 1-channel dimension (@Chandraveersingh1717, #307).
+* Fixed `draw_keypoints()` documentation and error message (@srishtiii28, #296).
 * Standardized dataset messages: download messages now include split information, success messages show image count and class count for consistency.
 * fix `model_fasterrcnn_*` did not provide boxes output normalized to image size, did not manage batches, fix performance of the `roi_align()` function (#284)
 * fix rf100 collection bounding-box now consider the correct native COCO format being 'xywh' (#272)
 * Remove `.getbatch` method from MNIST as it is providing inconsistent tensor dimensions with `.getitem` due
 to non-vectorized `transform_` operations (#264)
+* Added article for `draw_keypoints()` (@DerrickUnleashed #303)
+* Fix typos and align model documentation for `model_deeplabv3_*` and `model_convnext_*_detection()` to ensure consistency.(@DerrickUnleashed #302)
 
 # torchvision 0.8.0