Return real dataset length (#210)

4aff146c · Keon · Soumith Chintala · 7842c3ed · 4aff146c
Commit 4aff146c authored Jul 26, 2017 by Keon Committed by Soumith Chintala Jul 27, 2017
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

torchvision/datasets/mnist.py torchvision/datasets/mnist.py +2 -2

No files found.
--- a/torchvision/datasets/mnist.py
+++ b/torchvision/datasets/mnist.py
@@ -81,9 +81,9 @@ class MNIST(data.Dataset):
    def __len__(self):
        if self.train:
-            return 60000
+            return len(self.train_data)
        else:
-            return 10000
+            return len(self.test_data)
    def _check_exists(self):
        return os.path.exists(os.path.join(self.root, self.processed_folder, self.training_file)) and \