diff --git a/modules/varautoencoder_mednist.ipynb b/modules/varautoencoder_mednist.ipynb
index 3645b01331..8e47e90fad 100644
--- a/modules/varautoencoder_mednist.ipynb
+++ b/modules/varautoencoder_mednist.ipynb
@@ -358,9 +358,9 @@
     "    ]\n",
     ")\n",
     "\n",
-    "train_ds = CacheDataset(train_datadict[:10], transforms, num_workers=num_workers)\n",
+    "train_ds = CacheDataset(train_datadict, transforms, num_workers=num_workers)\n",
     "train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=num_workers)\n",
-    "test_ds = CacheDataset(test_datadict[:10], transforms, num_workers=num_workers)\n",
+    "test_ds = CacheDataset(test_datadict, transforms, num_workers=num_workers)\n",
     "test_loader = DataLoader(test_ds, batch_size=batch_size, shuffle=True, num_workers=num_workers)"
    ]
   },