diff --git a/abl/learning/__init__.py b/abl/learning/__init__.py
index d2dabfc..c3cfa0a 100644
--- a/abl/learning/__init__.py
+++ b/abl/learning/__init__.py
@@ -1,4 +1,5 @@
 from .abl_model import ABLModel
 from .basic_nn import BasicNN
+from .torch_dataset import *
 
 __all__ = ["ABLModel", "BasicNN"]
diff --git a/abl/learning/basic_nn.py b/abl/learning/basic_nn.py
index d47f87d..8bd92cb 100644
--- a/abl/learning/basic_nn.py
+++ b/abl/learning/basic_nn.py
@@ -8,7 +8,7 @@ import numpy
 import torch
 from torch.utils.data import DataLoader
 
-from ..dataset import ClassificationDataset, PredictionDataset
+from .torch_dataset import ClassificationDataset, PredictionDataset
 from ..utils.logger import print_log
 
 
diff --git a/abl/dataset/__init__.py b/abl/learning/torch_dataset/__init__.py
similarity index 100%
rename from abl/dataset/__init__.py
rename to abl/learning/torch_dataset/__init__.py
diff --git a/abl/dataset/classification_dataset.py b/abl/learning/torch_dataset/classification_dataset.py
similarity index 100%
rename from abl/dataset/classification_dataset.py
rename to abl/learning/torch_dataset/classification_dataset.py
diff --git a/abl/dataset/prediction_dataset.py b/abl/learning/torch_dataset/prediction_dataset.py
similarity index 100%
rename from abl/dataset/prediction_dataset.py
rename to abl/learning/torch_dataset/prediction_dataset.py
diff --git a/abl/dataset/regression_dataset.py b/abl/learning/torch_dataset/regression_dataset.py
similarity index 100%
rename from abl/dataset/regression_dataset.py
rename to abl/learning/torch_dataset/regression_dataset.py
diff --git a/docs/API/abl.dataset.rst b/docs/API/abl.dataset.rst
deleted file mode 100644
index 17e10e5..0000000
--- a/docs/API/abl.dataset.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-abl.dataset
-==================
-
-.. automodule:: abl.dataset
-   :members:
-   :undoc-members:
-   :show-inheritance:
\ No newline at end of file
diff --git a/docs/API/abl.learning.rst b/docs/API/abl.learning.rst
index f2c0f0c..72aff01 100644
--- a/docs/API/abl.learning.rst
+++ b/docs/API/abl.learning.rst
@@ -1,7 +1,23 @@
 abl.learning
 ==================
 
-.. automodule:: abl.learning
+Learning
+--------
+
+.. autoclass:: abl.learning.ABLModel
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: abl.learning.BasicNN
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+Torch Dataset
+-------------
+
+.. automodule:: abl.learning.torch_dataset
    :members:
    :undoc-members:
    :show-inheritance:
\ No newline at end of file
diff --git a/docs/index.rst b/docs/index.rst
index 6191a0b..b99a329 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -32,11 +32,10 @@
    :maxdepth: 1
    :caption: API
 
-   API/abl.dataset
+   API/abl.data
    API/abl.learning
    API/abl.reasoning
    API/abl.bridge
-   API/abl.data
    API/abl.utils
 
 .. toctree::
diff --git a/examples/hed/bridge.py b/examples/hed/bridge.py
index b7ca577..2852ca0 100644
--- a/examples/hed/bridge.py
+++ b/examples/hed/bridge.py
@@ -4,7 +4,7 @@ from collections import defaultdict
 import torch
 
 from abl.bridge import SimpleBridge
-from abl.dataset import RegressionDataset
+from abl.learning.torch_dataset import RegressionDataset
 from abl.data.evaluation import BaseMetric
 from abl.learning import ABLModel, BasicNN
 from abl.reasoning import Reasoner
diff --git a/examples/mnist_add/mnist_add.ipynb b/examples/mnist_add/mnist_add.ipynb
index 172e305..aed35b4 100644
--- a/examples/mnist_add/mnist_add.ipynb
+++ b/examples/mnist_add/mnist_add.ipynb
@@ -13,7 +13,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -42,7 +42,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -61,7 +61,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
@@ -110,7 +110,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [
     {
@@ -170,7 +170,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -198,7 +198,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [
     {
@@ -229,7 +229,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -245,7 +245,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [
     {
@@ -295,7 +295,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -319,7 +319,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 10,
    "metadata": {},
    "outputs": [
     {
@@ -352,7 +352,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -385,7 +385,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 12,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -404,7 +404,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 13,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -420,7 +420,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 14,
    "metadata": {},
    "outputs": [],
    "source": [