Repository: cornellius-gp/gpytorch
Pull request #676 (2019-05-04 00:14:02)

Add LinearMean
Seems like a reasonably simple parametric `Mean` for users to try out? 

gpleiss: Seems reasonable, once the tests are fixed!

Patch:
diff --git a/.travis.yml b/.travis.yml
index b1dddf7a..e8d3fecb 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -15,7 +15,7 @@ install:
       pip install torch_nightly -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html;
       python setup.py build develop;
     else
-      pip3 install https://download.pytorch.org/whl/cpu/torch-1.1.0-cp36-cp36m-linux_x86_64.whl;
+      pip3 install https://download.pytorch.org/whl/cpu/torch-1.0.1.post2-cp36-cp36m-linux_x86_64.whl;
     fi
 
 script:
diff --git a/examples/00_Basic_Usage/Hyperparameters.ipynb b/examples/00_Basic_Usage/Hyperparameters.ipynb
deleted file mode 100644
index 73436c9f..00000000
--- a/examples/00_Basic_Usage/Hyperparameters.ipynb
+++ /dev/null
@@ -1,513 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Hyperparameters in GPyTorch\n",
-    "\n",
-    "The purpose of this notebook is to explain how GP hyperparameters in GPyTorch work, how they are handled, what options are available for constraints and priors, and how things may differ from other packages.\n",
-    "\n",
-    "**Note:** This is a *basic* introduction to hyperparameters in GPyTorch. If you want to use GPyTorch hyperparameters with things like Pyro distributions, that will be covered in a less \"basic usage\" tutorial."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import math\n",
-    "import torch\n",
-    "import gpytorch\n",
-    "from matplotlib import pyplot as plt\n",
-    "\n",
-    "from IPython.display import Markdown, display\n",
-    "def printmd(string):\n",
-    "    display(Markdown(string))"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Defining an example model\n",
-    "\n",
-    "In the next cell, we define our simple exact GP from the <a href=\"../01_Simple_GP_Regression/Simple_GP_Regression.ipynb\">Simple GP Regression</a> tutorial. We'll be using this model to demonstrate certain aspects of hyperparameter creation."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "train_x = torch.linspace(0, 1, 100)\n",
-    "train_y = torch.sin(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * 0.2\n",
-    "\n",
-    "# We will use the simplest form of GP model, exact inference\n",
-    "class ExactGPModel(gpytorch.models.ExactGP):\n",
-    "    def __init__(self, train_x, train_y, likelihood):\n",
-    "        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)\n",
-    "        self.mean_module = gpytorch.means.ConstantMean()\n",
-    "        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())\n",
-    "    \n",
-    "    def forward(self, x):\n",
-    "        mean_x = self.mean_module(x)\n",
-    "        covar_x = self.covar_module(x)\n",
-    "        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)\n",
-    "\n",
-    "# initialize likelihood and model\n",
-    "likelihood = gpytorch.likelihoods.GaussianLikelihood()\n",
-    "model = ExactGPModel(train_x, train_y, likelihood)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Viewing model hyperparameters\n",
-    "\n",
-    "Let's take a look at the model parameters. By \"parameters\", here I mean explicitly objects of type `torch.nn.Parameter` that will have gradients filled in by autograd. To access these, there are two ways of doing this in torch. One way is to use `model.state_dict()`, which we demonstrate the use of for saving models <a href=\"Saving_and_Loading_Models.ipynb\">here</a>.\n",
-    "\n",
-    "In the next cell we demonstrate another way to do this, by looping over the `model.named_parameters()` generator:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Parameter name: likelihood.noise_covar.raw_noise           value = 0.0\n",
-      "Parameter name: mean_module.constant                       value = 0.0\n",
-      "Parameter name: covar_module.raw_outputscale               value = 0.0\n",
-      "Parameter name: covar_module.base_kernel.raw_lengthscale   value = 0.0\n"
-     ]
-    }
-   ],
-   "source": [
-    "for param_name, param in model.named_parameters():\n",
-    "    print(f'Parameter name: {param_name:42} value = {param.item()}')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Raw vs Actual Parameters\n",
-    "\n",
-    "The most important thing to note here is that the actual learned parameters of the model are things like `raw_noise`, `raw_outputscale`, `raw_lengthscale`, etc. The reason for this is that these parameters **must be positive**. This brings us to our next topic for parameters: constraints, and the difference between *raw* parameters and *actual* parameters.\n",
-    "\n",
-    "In order to enforce positiveness and other constraints for hyperparameters, GPyTorch has **raw** parameters (e.g., `model.covar_module.raw_outputscale`) that are transformed to actual values via some constraint. Let's take a look at the raw outputscale, its constraint, and the final value:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "raw_outputscale,  Parameter containing:\n",
-      "tensor(0., requires_grad=True)\n",
-      "\n",
-      "raw_outputscale_constraint1 Positive()\n"
-     ]
-    },
-    {
-     "data": {
-      "text/markdown": [
-       "\n",
-       "\n",
-       "**Printing all model constraints...**\n"
-      ],
-      "text/plain": [
-       "<IPython.core.display.Markdown object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Constraint name: likelihood.noise_covar.raw_noise_constraint             constraint = GreaterThan(1.000E-04)\n",
-      "Constraint name: covar_module.raw_outputscale_constraint                 constraint = Positive()\n",
-      "Constraint name: covar_module.base_kernel.raw_lengthscale_constraint     constraint = Positive()\n"
-     ]
-    },
-    {
-     "data": {
-      "text/markdown": [
-       "\n",
-       "**Getting raw outputscale constraint from model...**"
-      ],
-      "text/plain": [
-       "<IPython.core.display.Markdown object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Positive()\n"
-     ]
-    },
-    {
-     "data": {
-      "text/markdown": [
-       "\n",
-       "**Getting raw outputscale constraint from model.covar_module...**"
-      ],
-      "text/plain": [
-       "<IPython.core.display.Markdown object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Positive()\n"
-     ]
-    }
-   ],
-   "source": [
-    "raw_outputscale = model.covar_module.raw_outputscale\n",
-    "print('raw_outputscale, ', raw_outputscale)\n",
-    "\n",
-    "# Three ways of accessing the raw outputscale constraint\n",
-    "print('\\nraw_outputscale_constraint1', model.covar_module.raw_outputscale_constraint)\n",
-    "\n",
-    "printmd('\\n\\n**Printing all model constraints...**\\n')\n",
-    "for constraint_name, constraint in model.named_constraints():\n",
-    "    print(f'Constraint name: {constraint_name:55} constraint = {constraint}')\n",
-    "\n",
-    "printmd('\\n**Getting raw outputscale constraint from model...**')\n",
-    "print(model.constraint_for_parameter_name(\"covar_module.raw_outputscale\"))\n",
-    "\n",
-    "\n",
-    "printmd('\\n**Getting raw outputscale constraint from model.covar_module...**')\n",
-    "print(model.covar_module.constraint_for_parameter_name(\"raw_outputscale\"))"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## How do constraints work?\n",
-    "\n",
-    "Constraints define `transform` and `inverse_transform` methods that turn raw parameters in to real ones. For a positive constraint, we expect the **transformed** values to always be positive. Let's see:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Transformed outputscale tensor(0.6931, grad_fn=<SoftplusBackward>)\n",
-      "tensor(0., grad_fn=<LogBackward>)\n",
-      "True\n",
-      "Transform a bunch of negative tensors:  tensor([0.3133, 0.1269, 0.0486])\n"
-     ]
-    }
-   ],
-   "source": [
-    "raw_outputscale = model.covar_module.raw_outputscale\n",
-    "constraint = model.covar_module.raw_outputscale_constraint\n",
-    "\n",
-    "print('Transformed outputscale', constraint.transform(raw_outputscale))\n",
-    "print(constraint.inverse_transform(constraint.transform(raw_outputscale)))\n",
-    "print(torch.equal(constraint.inverse_transform(constraint.transform(raw_outputscale)), raw_outputscale))\n",
-    "\n",
-    "print('Transform a bunch of negative tensors: ', constraint.transform(torch.tensor([-1., -2., -3.])))"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Convenience Getters/Setters for Transformed Values\n",
-    "\n",
-    "Because dealing with raw parameter values is annoying (e.g., we might know what a noise variance of 0.01 means, but maybe not a `raw_noise` of `-2.791`), virtually all built in GPyTorch modules that define raw parameters define convenience getters and setters for dealing with transformed values directly.\n",
-    "\n",
-    "In the next cells, we demonstrate the \"inconvenient way\" and the \"convenient\" way of getting and setting the outputscale."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Actual outputscale: 0.6931471824645996\n",
-      "Actual outputscale after setting: 2.0\n"
-     ]
-    }
-   ],
-   "source": [
-    "# Recreate model to reset outputscale\n",
-    "model = ExactGPModel(train_x, train_y, likelihood)\n",
-    "\n",
-    "## Inconvenient way of getting true outputscale\n",
-    "raw_outputscale = model.covar_module.raw_outputscale\n",
-    "constraint = model.covar_module.raw_outputscale_constraint\n",
-    "outputscale = constraint.transform(raw_outputscale)\n",
-    "print(f'Actual outputscale: {outputscale.item()}')\n",
-    "\n",
-    "## Inconvenient way of setting true outputscale\n",
-    "model.covar_module.raw_outputscale.data.fill_(constraint.inverse_transform(torch.tensor(2.)))\n",
-    "raw_outputscale = model.covar_module.raw_outputscale\n",
-    "outputscale = constraint.transform(raw_outputscale)\n",
-    "print(f'Actual outputscale after setting: {outputscale.item()}')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Ouch, that is ugly! Fortunately, there is a better way:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 7,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Actual outputscale: 0.6931471824645996\n",
-      "Actual outputscale after setting: 2.0\n"
-     ]
-    }
-   ],
-   "source": [
-    "# Recreate model to reset outputscale\n",
-    "model = ExactGPModel(train_x, train_y, likelihood)\n",
-    "\n",
-    "# Convenient way of getting true outputscale\n",
-    "print(f'Actual outputscale: {model.covar_module.outputscale}')\n",
-    "\n",
-    "# Convenient way of setting true outputscale\n",
-    "model.covar_module.outputscale = 2.\n",
-    "print(f'Actual outputscale after setting: {model.covar_module.outputscale}')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Changing Parameter Constraints\n",
-    "\n",
-    "If we look at the actual noise of the model, GPyTorch defines a default lower bound of `1e-4` for the noise variance:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Actual noise value: tensor([0.6932], grad_fn=<AddBackward0>)\n",
-      "Noise constraint: GreaterThan(1.000E-04)\n"
-     ]
-    }
-   ],
-   "source": [
-    "print(f'Actual noise value: {likelihood.noise}')\n",
-    "print(f'Noise constraint: {likelihood.noise_covar.raw_noise_constraint}')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "We can change the noise constraint either on the fly or when the likelihood is created:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Noise constraint: GreaterThan(1.000E-03)\n",
-      "Noise constraint: Positive()\n"
-     ]
-    }
-   ],
-   "source": [
-    "likelihood = gpytorch.likelihoods.GaussianLikelihood(noise_constraint=gpytorch.constraints.GreaterThan(1e-3))\n",
-    "print(f'Noise constraint: {likelihood.noise_covar.raw_noise_constraint}')\n",
-    "\n",
-    "# Changing the constraint after the module has been created\n",
-    "likelihood.noise_covar.register_constraint(\"raw_noise\", gpytorch.constraints.Positive())\n",
-    "print(f'Noise constraint: {likelihood.noise_covar.raw_noise_constraint}')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Priors\n",
-    "\n",
-    "In GPyTorch, priors are things you register to the model that act on any arbitrary function of any parameter. Like constraints, these can usually be defined either when you create an object (like a Kernel or Likelihood), or set afterwards on the fly.\n",
-    "\n",
-    "Here are some examples:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 10,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Registers a prior on the sqrt of the noise parameter \n",
-    "# (e.g., a prior for the noise standard deviation instead of variance)\n",
-    "likelihood.noise_covar.register_prior(\n",
-    "    \"noise_std_prior\",\n",
-    "    gpytorch.priors.NormalPrior(0, 1),\n",
-    "    lambda: likelihood.noise.sqrt()\n",
-    ")\n",
-    "\n",
-    "# Create a GaussianLikelihood with a normal prior for the noise\n",
-    "likelihood = gpytorch.likelihoods.GaussianLikelihood(\n",
-    "    noise_constraint=gpytorch.constraints.GreaterThan(1e-3),\n",
-    "    noise_prior=gpytorch.priors.NormalPrior(0, 1)\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Putting it Together\n",
-    "\n",
-    "In the next cell, we augment our `ExactGP` definition to place several priors over hyperparameters and tighter constraints when creating the model."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 11,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# We will use the simplest form of GP model, exact inference\n",
-    "class FancyGPWithPriors(gpytorch.models.ExactGP):\n",
-    "    def __init__(self, train_x, train_y, likelihood):\n",
-    "        super(FancyGPWithPriors, self).__init__(train_x, train_y, likelihood)\n",
-    "        self.mean_module = gpytorch.means.ConstantMean()\n",
-    "        \n",
-    "        lengthscale_prior = gpytorch.priors.GammaPrior(3.0, 6.0)\n",
-    "        outputscale_prior = gpytorch.priors.GammaPrior(2.0, 0.15)\n",
-    "        \n",
-    "        self.covar_module = gpytorch.kernels.ScaleKernel(\n",
-    "            gpytorch.kernels.RBFKernel(\n",
-    "                lengthscale_prior=lengthscale_prior,\n",
-    "            ),\n",
-    "            outputscale_prior=outputscale_prior\n",
-    "        )\n",
-    "        \n",
-    "        # Initialize lengthscale and outputscale to mean of priors\n",
-    "        self.covar_module.base_kernel.lengthscale = lengthscale_prior.mean\n",
-    "        self.covar_module.outputscale = outputscale_prior.mean\n",
-    "    \n",
-    "    def forward(self, x):\n",
-    "        mean_x = self.mean_module(x)\n",
-    "        covar_x = self.covar_module(x)\n",
-    "        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)\n",
-    "\n",
-    "likelihood = gpytorch.likelihoods.GaussianLikelihood(\n",
-    "    noise_constraint=gpytorch.constraints.GreaterThan(1e-2),\n",
-    ")\n",
-    "\n",
-    "model = FancyGPWithPriors(train_x, train_y, likelihood)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Initializing hyperparameters in One Call\n",
-    "\n",
-    "For convenience, GPyTorch modules also define an `initialize` method that allow you to update a full dictionary of parameters on submodules. For example:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 33,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "1.0000001192092896 0.4999999701976776 2.0\n"
-     ]
-    }
-   ],
-   "source": [
-    "hypers = {\n",
-    "    'likelihood.noise_covar.noise': torch.tensor(1.),\n",
-    "    'covar_module.base_kernel.lengthscale': torch.tensor(0.5),\n",
-    "    'covar_module.outputscale': torch.tensor(2.),\n",
-    "}\n",
-    "\n",
-    "model.initialize(**hypers)\n",
-    "print(\n",
-    "    model.likelihood.noise_covar.noise.item(),\n",
-    "    model.covar_module.base_kernel.lengthscale.item(),\n",
-    "    model.covar_module.outputscale.item()\n",
-    ")"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.7.1"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}
diff --git a/examples/00_Basic_Usage/README.md b/examples/00_Basic_Usage/README.md
deleted file mode 100644
index 54e1e993..00000000
--- a/examples/00_Basic_Usage/README.md
+++ /dev/null
@@ -1,10 +0,0 @@
-# Basic Usage
-
-This folder contains notebooks for basic usage of the package, e.g. things like dealing with hyperparameters,
-parameter constraints and priors, and saving and loading models. Before checking these out, you may want to see our
-simple GP regression notebooks that detail the anatomy of a GPyTorch model [here](../01_Simple_GP_Regression).
-
-- Check out our [Tutorial on Hyperparameters](./Hyperparameters.ipynb) for information on things like raw versus actual
-  parameters, constraints, priors and more.
-- The [Saving and loading models](./Saving_and_Loading_Models.ipynb) notebook details how to save and load GPyTorch models
-  on disk.
\ No newline at end of file
diff --git a/examples/00_Basic_Usage/Saving_and_Loading_Models.ipynb b/examples/00_Basic_Usage/Saving_and_Loading_Models.ipynb
deleted file mode 100644
index 8346d062..00000000
--- a/examples/00_Basic_Usage/Saving_and_Loading_Models.ipynb
+++ /dev/null
@@ -1,371 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Saving and Loading Models\n",
-    "\n",
-    "In this bite-sized notebook, we'll go over how to save and load models. In general, the process is the same as for any PyTorch module."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import math\n",
-    "import torch\n",
-    "import gpytorch\n",
-    "from matplotlib import pyplot as plt"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Saving a Simple Model\n",
-    "\n",
-    "First, we define a GP Model that we'd like to save. The model used below is the same as the model from our\n",
-    "<a href=\"../01_Simple_GP_Regression/Simple_GP_Regression.ipynb\">Simple GP Regression</a> tutorial."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "train_x = torch.linspace(0, 1, 100)\n",
-    "train_y = torch.sin(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * 0.2"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# We will use the simplest form of GP model, exact inference\n",
-    "class ExactGPModel(gpytorch.models.ExactGP):\n",
-    "    def __init__(self, train_x, train_y, likelihood):\n",
-    "        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)\n",
-    "        self.mean_module = gpytorch.means.ConstantMean()\n",
-    "        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())\n",
-    "    \n",
-    "    def forward(self, x):\n",
-    "        mean_x = self.mean_module(x)\n",
-    "        covar_x = self.covar_module(x)\n",
-    "        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)\n",
-    "\n",
-    "# initialize likelihood and model\n",
-    "likelihood = gpytorch.likelihoods.GaussianLikelihood()\n",
-    "model = ExactGPModel(train_x, train_y, likelihood)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Change Model State\n",
-    "\n",
-    "To demonstrate model saving, we change the hyperparameters from the default values below. For more information on what is happening here, see our tutorial notebook on <a href=\"Hyperparameters.ipynb\">Initializing Hyperparameters</a>."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model.covar_module.outputscale = 1.2\n",
-    "model.covar_module.base_kernel.lengthscale = 2.2"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Getting Model State\n",
-    "\n",
-    "To get the full state of a GPyTorch model, simply call `state_dict` as you would on any PyTorch model. Note that the state dict contains **raw** parameter values. This is because these are the actual `torch.nn.Parameters` that are learned in GPyTorch. Again see our notebook on hyperparamters for more information on this."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 11,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "OrderedDict([('likelihood.noise_covar.raw_noise', tensor([0.])),\n",
-       "             ('mean_module.constant', tensor([0.])),\n",
-       "             ('covar_module.raw_outputscale', tensor(0.8416)),\n",
-       "             ('covar_module.base_kernel.raw_lengthscale', tensor([[2.0826]]))])"
-      ]
-     },
-     "execution_count": 11,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "model.state_dict()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Saving Model State\n",
-    "\n",
-    "The state dictionary above represents all traininable parameters for the model. Therefore, we can save this to a file as follows:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "torch.save(model.state_dict(), 'model_state.pth')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Loading Model State\n",
-    "\n",
-    "Next, we load this state in to a new model and demonstrate that the parameters were updated correctly."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 13,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "IncompatibleKeys(missing_keys=[], unexpected_keys=[])"
-      ]
-     },
-     "execution_count": 13,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "state_dict = torch.load('model_state.pth')\n",
-    "model = ExactGPModel(train_x, train_y, likelihood)  # Create a new GP model\n",
-    "\n",
-    "model.load_state_dict(state_dict)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 15,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "OrderedDict([('likelihood.noise_covar.raw_noise', tensor([0.])),\n",
-       "             ('mean_module.constant', tensor([0.])),\n",
-       "             ('covar_module.raw_outputscale', tensor(0.8416)),\n",
-       "             ('covar_module.base_kernel.raw_lengthscale', tensor([[2.0826]]))])"
-      ]
-     },
-     "execution_count": 15,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "model.state_dict()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# A More Complex Example\n",
-    "\n",
-    "Next we demonstrate this same principle on a more complex exact GP where we have a simple feed forward neural network feature extractor as part of the model.\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 22,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "class GPWithNNFeatureExtractor(gpytorch.models.ExactGP):\n",
-    "    def __init__(self, train_x, train_y, likelihood):\n",
-    "        super(GPWithNNFeatureExtractor, self).__init__(train_x, train_y, likelihood)\n",
-    "        self.mean_module = gpytorch.means.ConstantMean()\n",
-    "        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())\n",
-    "        \n",
-    "        self.feature_extractor = torch.nn.Sequential(\n",
-    "            torch.nn.Linear(1, 2),\n",
-    "            torch.nn.BatchNorm1d(2),\n",
-    "            torch.nn.ReLU(),\n",
-    "            torch.nn.Linear(2, 2),\n",
-    "            torch.nn.BatchNorm1d(2),\n",
-    "            torch.nn.ReLU(),\n",
-    "        )\n",
-    "    \n",
-    "    def forward(self, x):\n",
-    "        x = self.feature_extractor(x)\n",
-    "        mean_x = self.mean_module(x)\n",
-    "        covar_x = self.covar_module(x)\n",
-    "        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)\n",
-    "\n",
-    "# initialize likelihood and model\n",
-    "likelihood = gpytorch.likelihoods.GaussianLikelihood()\n",
-    "model = GPWithNNFeatureExtractor(train_x, train_y, likelihood)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Getting Model State\n",
-    "\n",
-    "In the next cell, we once again print the model state via `model.state_dict()`. As you can see, the state is substantially more complex, as the model now includes our neural network parameters. Nevertheless, saving and loading is straight forward."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 23,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "OrderedDict([('likelihood.noise_covar.raw_noise', tensor([0.])),\n",
-       "             ('mean_module.constant', tensor([0.])),\n",
-       "             ('covar_module.raw_outputscale', tensor(0.)),\n",
-       "             ('covar_module.base_kernel.raw_lengthscale', tensor([[0.]])),\n",
-       "             ('feature_extractor.0.weight', tensor([[-0.9135],\n",
-       "                      [-0.5942]])),\n",
-       "             ('feature_extractor.0.bias', tensor([ 0.9119, -0.0663])),\n",
-       "             ('feature_extractor.1.weight', tensor([0.2263, 0.2209])),\n",
-       "             ('feature_extractor.1.bias', tensor([0., 0.])),\n",
-       "             ('feature_extractor.1.running_mean', tensor([0., 0.])),\n",
-       "             ('feature_extractor.1.running_var', tensor([1., 1.])),\n",
-       "             ('feature_extractor.1.num_batches_tracked', tensor(0)),\n",
-       "             ('feature_extractor.3.weight', tensor([[-0.6375, -0.6466],\n",
-       "                      [-0.0563, -0.4695]])),\n",
-       "             ('feature_extractor.3.bias', tensor([-0.1247,  0.0803])),\n",
-       "             ('feature_extractor.4.weight', tensor([0.0466, 0.7248])),\n",
-       "             ('feature_extractor.4.bias', tensor([0., 0.])),\n",
-       "             ('feature_extractor.4.running_mean', tensor([0., 0.])),\n",
-       "             ('feature_extractor.4.running_var', tensor([1., 1.])),\n",
-       "             ('feature_extractor.4.num_batches_tracked', tensor(0))])"
-      ]
-     },
-     "execution_count": 23,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "model.state_dict()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 24,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "IncompatibleKeys(missing_keys=[], unexpected_keys=[])"
-      ]
-     },
-     "execution_count": 24,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "torch.save(model.state_dict(), 'my_gp_with_nn_model.pth')\n",
-    "state_dict = torch.load('my_gp_with_nn_model.pth')\n",
-    "model = GPWithNNFeatureExtractor(train_x, train_y, likelihood)\n",
-    "model.load_state_dict(state_dict)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 25,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "OrderedDict([('likelihood.noise_covar.raw_noise', tensor([0.])),\n",
-       "             ('mean_module.constant', tensor([0.])),\n",
-       "             ('covar_module.raw_outputscale', tensor(0.)),\n",
-       "             ('covar_module.base_kernel.raw_lengthscale', tensor([[0.]])),\n",
-       "             ('feature_extractor.0.weight', tensor([[-0.9135],\n",
-       "                      [-0.5942]])),\n",
-       "             ('feature_extractor.0.bias', tensor([ 0.9119, -0.0663])),\n",
-       "             ('feature_extractor.1.weight', tensor([0.2263, 0.2209])),\n",
-       "             ('feature_extractor.1.bias', tensor([0., 0.])),\n",
-       "             ('feature_extractor.1.running_mean', tensor([0., 0.])),\n",
-       "             ('feature_extractor.1.running_var', tensor([1., 1.])),\n",
-       "             ('feature_extractor.1.num_batches_tracked', tensor(0)),\n",
-       "             ('feature_extractor.3.weight', tensor([[-0.6375, -0.6466],\n",
-       "                      [-0.0563, -0.4695]])),\n",
-       "             ('feature_extractor.3.bias', tensor([-0.1247,  0.0803])),\n",
-       "             ('feature_extractor.4.weight', tensor([0.0466, 0.7248])),\n",
-       "             ('feature_extractor.4.bias', tensor([0., 0.])),\n",
-       "             ('feature_extractor.4.running_mean', tensor([0., 0.])),\n",
-       "             ('feature_extractor.4.running_var', tensor([1., 1.])),\n",
-       "             ('feature_extractor.4.num_batches_tracked', tensor(0))])"
-      ]
-     },
-     "execution_count": 25,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "model.state_dict()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.7.1"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}
diff --git a/examples/00_Basic_Usage/index.rst b/examples/00_Basic_Usage/index.rst
deleted file mode 100644
index 7a814057..00000000
--- a/examples/00_Basic_Usage/index.rst
+++ /dev/null
@@ -1,8 +0,0 @@
-.. mdinclude:: README.md
-
-.. toctree::
-   :maxdepth: 1
-   :hidden:
-
-   Hyperparameters.ipynb
-   Saving_and_Loading_Models.ipynb
diff --git a/examples/01_Simple_GP_Regression/Simple_MultiGPU_GP_Regression.ipynb b/examples/01_Simple_GP_Regression/Simple_MultiGPU_GP_Regression.ipynb
index 2ee0cef7..4a3959bf 100644
--- a/examples/01_Simple_GP_Regression/Simple_MultiGPU_GP_Regression.ipynb
+++ b/examples/01_Simple_GP_Regression/Simple_MultiGPU_GP_Regression.ipynb
@@ -16,7 +16,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 2,
    "metadata": {
     "scrolled": true
    },
@@ -45,7 +45,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -72,7 +72,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -116,7 +116,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [
     {
@@ -143,7 +143,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -222,19 +222,37 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Number of devices: 2 -- Kernel partition size: 0\n",
-      "RuntimeError: CUDA out of memory. Tried to allocate 2.49 GiB (GPU 1; 10.73 GiB total capacity; 7.48 GiB already allocated; 2.46 GiB free; 21.49 MiB cached)\n",
+      "Number of devices: 2 -- Kernel partition size: 0\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/jrg365/anaconda3/lib/python3.6/site-packages/torch/nn/parallel/data_parallel.py:26: UserWarning: \n",
+      "    There is an imbalance between your GPUs. You may want to exclude GPU 1 which\n",
+      "    has less than 75% of the memory or cores of GPU 0. You can do so by setting\n",
+      "    the device_ids argument to DataParallel, or by setting the CUDA_VISIBLE_DEVICES\n",
+      "    environment variable.\n",
+      "  warnings.warn(imbalance_warn.format(device_ids[min_pos], device_ids[max_pos]))\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "RuntimeError: CUDA out of memory. Tried to allocate 2.49 GiB (GPU 0; 10.73 GiB total capacity; 4.99 GiB already allocated; 2.39 GiB free; 22.42 MiB cached)\n",
       "Number of devices: 2 -- Kernel partition size: 18292\n",
-      "RuntimeError: CUDA out of memory. Tried to allocate 1.25 GiB (GPU 0; 10.73 GiB total capacity; 6.37 GiB already allocated; 448.94 MiB free; 1.30 GiB cached)\n",
+      "RuntimeError: CUDA out of memory. Tried to allocate 1.25 GiB (GPU 1; 7.93 GiB total capacity; 6.23 GiB already allocated; 1.19 GiB free; 28.49 MiB cached)\n",
       "Number of devices: 2 -- Kernel partition size: 9146\n",
-      "Iter 1/1 - Loss: 0.893   lengthscale: 0.486   noise: 0.248\n",
+      "Iter 1/1 - Loss: 0.947   lengthscale: 0.583   noise: 0.428\n",
       "Finished training on 36584 data points using 2 GPUs.\n"
      ]
     }
@@ -292,19 +310,27 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 23,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Iter 1/20 - Loss: 0.897   lengthscale: 0.486   noise: 0.253\n",
-      "Iter 2/20 - Loss: 0.892   lengthscale: 0.354   noise: 0.144\n",
-      "Iter 3/20 - Loss: 0.859   lengthscale: 0.305   noise: 0.125\n",
-      "Iter 4/20 - Loss: 0.859   lengthscale: 0.292   noise: 0.116\n",
-      "Iter 5/20 - Loss: 0.862   lengthscale: 0.292   noise: 0.116\n",
-      "Convergence reached!\n",
+      "Iter 1/20 - Loss: 1.030   lengthscale: 0.671   noise: 0.570\n",
+      "Iter 2/20 - Loss: 0.995   lengthscale: 0.606   noise: 0.476\n",
+      "Iter 3/20 - Loss: 0.969   lengthscale: 0.534   noise: 0.410\n",
+      "Iter 4/20 - Loss: 0.949   lengthscale: 0.473   noise: 0.360\n",
+      "Iter 5/20 - Loss: 0.935   lengthscale: 0.442   noise: 0.325\n",
+      "Iter 6/20 - Loss: 0.927   lengthscale: 0.418   noise: 0.302\n",
+      "Iter 7/20 - Loss: 0.912   lengthscale: 0.383   noise: 0.267\n",
+      "Iter 8/20 - Loss: 0.907   lengthscale: 0.358   noise: 0.243\n",
+      "Iter 9/20 - Loss: 0.904   lengthscale: 0.342   noise: 0.225\n",
+      "Iter 10/20 - Loss: 0.898   lengthscale: 0.326   noise: 0.212\n",
+      "Iter 11/20 - Loss: 0.897   lengthscale: 0.320   noise: 0.205\n",
+      "Iter 12/20 - Loss: 0.895   lengthscale: 0.318   noise: 0.200\n",
+      "Iter 13/20 - Loss: 0.895   lengthscale: 0.313   noise: 0.193\n",
+      "Iter 14/20 - Loss: 0.896   lengthscale: 0.313   noise: 0.193\n",
       "Finished training on 36584 data points using 2 GPUs.\n"
      ]
     }
@@ -312,8 +338,8 @@
    "source": [
     "model, likelihood = train(train_x, train_y,\n",
     "                          n_devices=n_devices, output_device=output_device,\n",
-    "                          checkpoint_size=10000,\n",
-    "                          preconditioner_size=100,\n",
+    "                          checkpoint_size=checkpoint_size,\n",
+    "                          preconditioner_size=preconditioner_size,\n",
     "                          n_training_iter=20)"
    ]
   },
@@ -326,7 +352,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 24,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -334,10 +360,8 @@
     "model.eval()\n",
     "likelihood.eval()\n",
     "\n",
-    "with torch.no_grad(), gpytorch.settings.fast_pred_var(), gpytorch.beta_features.checkpoint_kernel(1000):\n",
-    "    # Make predictions on a small number of test points to get the test time caches computed\n",
-    "    latent_pred = model(test_x[:2, :])\n",
-    "    del latent_pred  # We don't care about these predictions, we really just want the caches."
+    "with torch.no_grad(), gpytorch.settings.fast_pred_var():\n",
+    "    latent_pred = model(test_x)"
    ]
   },
   {
@@ -349,34 +373,27 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 25,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "CPU times: user 1.1 s, sys: 728 ms, total: 1.83 s\n",
-      "Wall time: 1.88 s\n",
-      "Test RMSE: 0.551821768283844\n"
+      "CPU times: user 53 ms, sys: 12.3 ms, total: 65.3 ms\n",
+      "Wall time: 63 ms\n",
+      "Test RMSE: 0.5422027707099915\n"
      ]
     }
    ],
    "source": [
-    "with torch.no_grad(), gpytorch.settings.fast_pred_var(), gpytorch.beta_features.checkpoint_kernel(1000):\n",
+    "with torch.no_grad(), gpytorch.settings.fast_pred_var():\n",
     "    %time latent_pred = model(test_x)\n",
     "    \n",
     "test_rmse = torch.sqrt(torch.mean(torch.pow(latent_pred.mean - test_y, 2)))\n",
     "print(f\"Test RMSE: {test_rmse.item()}\")"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": 10,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -402,7 +419,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.1"
+   "version": "3.6.7"
   }
  },
  "nbformat": 4,
diff --git a/gpytorch/functions/_inv_quad_log_det.py b/gpytorch/functions/_inv_quad_log_det.py
index 850986e1..f293a7db 100644
--- a/gpytorch/functions/_inv_quad_log_det.py
+++ b/gpytorch/functions/_inv_quad_log_det.py
@@ -60,7 +60,7 @@ class InvQuadLogDet(Function):
         # Get closure for matmul
         lazy_tsr = ctx.representation_tree(*matrix_args)
         with torch.no_grad():
-            preconditioner, precond_lt, logdet_correction = lazy_tsr._preconditioner()
+            preconditioner, precond_lt, logdet_correction = lazy_tsr.detach()._preconditioner()
 
         ctx.preconditioner = preconditioner
 
diff --git a/gpytorch/kernels/spectral_mixture_kernel.py b/gpytorch/kernels/spectral_mixture_kernel.py
index df875a35..bc7f344b 100644
--- a/gpytorch/kernels/spectral_mixture_kernel.py
+++ b/gpytorch/kernels/spectral_mixture_kernel.py
@@ -18,7 +18,7 @@ class SpectralMixtureKernel(Kernel):
     .. note::
 
         Unlike other kernels,
-        * :attr:`ard_num_dims` **must equal** the number of dimensions of the data
+        * :attr:`ard_num_dums` **must equal** the number of dimensions of the data
         * :attr:`batch_shape` **must equal** the batch size of the data (torch.Size([1]) if the data is not batched)
         * :attr:`batch_shape` **cannot** contain more than one batch dimension.
         * This kernel should not be combined with a :class:`gpytorch.kernels.ScaleKernel`.
@@ -50,12 +50,12 @@ class SpectralMixtureKernel(Kernel):
     Example:
         >>> # Non-batch
         >>> x = torch.randn(10, 5)
-        >>> covar_module = gpytorch.kernels.SpectralMixtureKernel(num_mixtures=4, ard_num_dims=5)
+        >>> covar_module = gpytorch.kernels.SpectralMixtureKernel(num_mixtures=4, ard_dum_dims=5)
         >>> covar = covar_module(x)  # Output: LazyVariable of size (10 x 10)
         >>>
         >>> # Batch
         >>> batch_x = torch.randn(2, 10, 5)
-        >>> covar_module = gpytorch.kernels.SpectralMixtureKernel(num_mixtures=4, batch_size=2, ard_num_dims=5)
+        >>> covar_module = gpytorch.kernels.SpectralMixtureKernel(num_mixtures=4, batch_size=2, ard_dum_dims=5)
         >>> covar = covar_module(x)  # Output: LazyVariable of size (10 x 10)
 
 
diff --git a/gpytorch/lazy/added_diag_lazy_tensor.py b/gpytorch/lazy/added_diag_lazy_tensor.py
index 1e75bd35..ccdc32c4 100644
--- a/gpytorch/lazy/added_diag_lazy_tensor.py
+++ b/gpytorch/lazy/added_diag_lazy_tensor.py
@@ -67,23 +67,26 @@ class AddedDiagLazyTensor(SumLazyTensor):
                 *batch_shape, n, k = self._piv_chol_self.shape
                 self._noise = self._diag_tensor.diag().unsqueeze(-1)
 
-                self.constant_diag = torch.equal(self._noise, self._noise[0] * torch.ones_like(self._noise))
+                constant_diag = torch.equal(self._noise, self._noise[0] * torch.ones_like(self._noise))
 
                 eye = torch.eye(k, dtype=self._piv_chol_self.dtype, device=self._piv_chol_self.device)
 
-                if self.constant_diag:
+                if constant_diag:
                     # We can factor out the noise for for both QR and solves.
-                    self.noise_constant = self._noise[0].squeeze()
+                    noise_constant = self._noise[0].squeeze()
                     self._q_cache, self._r_cache = torch.qr(
-                        torch.cat((self._piv_chol_self, self.noise_constant.sqrt() * eye))
+                        torch.cat((self._piv_chol_self, noise_constant.sqrt() * eye))
                     )
                     self._q_cache = self._q_cache[:n, :]
 
                     # Use the matrix determinant lemma for the logdet, using the fact that R'R = L_k'L_k + s*I
                     logdet = self._r_cache.diagonal(dim1=-1, dim2=-2).abs().log().sum(-1).mul(2)
-                    logdet = logdet + (n - k) * self.noise_constant.log()
+                    logdet = logdet + (n - k) * noise_constant.log()
                     self._precond_logdet_cache = logdet.view(*batch_shape) if len(batch_shape) else logdet.squeeze()
 
+                    def precondition_closure(rhs: torch.Tensor):
+                        return (1 / noise_constant) * (rhs - self._q_cache.matmul(self._q_cache.t().matmul(rhs)))
+
                 else:
                     # With non-constant diagonals, we cant factor out the noise as easily
                     self._q_cache, self._r_cache = torch.qr(torch.cat((self._piv_chol_self / self._noise.sqrt(), eye)))
@@ -94,6 +97,9 @@ class AddedDiagLazyTensor(SumLazyTensor):
                     ).log().sum([-1, -2])
                     self._precond_logdet_cache = logdet.view(*batch_shape) if len(batch_shape) else logdet.squeeze()
 
+                    def precondition_closure(rhs: torch.Tensor):
+                        return (rhs / self._noise) - self._q_cache.matmul(self._q_cache.t().matmul(rhs))
+
             else:
                 self._woodbury_cache, self._inv_scale, self._precond_logdet_cache = woodbury.woodbury_factor(
                     self._piv_chol_self, self._piv_chol_self, self._diag_tensor.diag(), logdet=True
@@ -101,27 +107,19 @@ class AddedDiagLazyTensor(SumLazyTensor):
                 self._scaled_inv_diag = self._inv_scale / self._diag_tensor.diag().unsqueeze(-1)
                 self._scaled_inv_diag_piv_chol_self = self._piv_chol_self * self._scaled_inv_diag
 
-            self.preconditioner_lt = PsdSumLazyTensor(RootLazyTensor(self._piv_chol_self), self._diag_tensor)
-
-        # NOTE to future self:
-        # We cannot memoize this precondition closure
-        # It causes a memory leak otherwise
-        def precondition_closure(tensor):
-            if hasattr(self, "_q_cache"):
-                if self.constant_diag:
-                    return (1 / self.noise_constant) * (tensor - self._q_cache.matmul(self._q_cache.t().matmul(tensor)))
+                # preconditioner
+                def precondition_closure(tensor):
+                    res = woodbury.woodbury_solve(
+                        tensor,
+                        self._scaled_inv_diag_piv_chol_self,
+                        self._woodbury_cache,
+                        self._scaled_inv_diag,
+                        self._inv_scale,
+                    )
+                    return res
 
-                else:
-                    return (tensor / self._noise) - self._q_cache.matmul(self._q_cache.t().matmul(tensor))
+            self.preconditioner_lt = PsdSumLazyTensor(RootLazyTensor(self._piv_chol_self), self._diag_tensor)
 
-            else:
-                res = woodbury.woodbury_solve(
-                    tensor,
-                    self._scaled_inv_diag_piv_chol_self,
-                    self._woodbury_cache,
-                    self._scaled_inv_diag,
-                    self._inv_scale,
-                )
-                return res
+            self.precondition_closure = precondition_closure
 
-        return precondition_closure, self.preconditioner_lt, self._precond_logdet_cache
+        return self.precondition_closure, self.preconditioner_lt, self._precond_logdet_cache
diff --git a/gpytorch/means/__init__.py b/gpytorch/means/__init__.py
index 3c9a7eff..0bbc86eb 100644
--- a/gpytorch/means/__init__.py
+++ b/gpytorch/means/__init__.py
@@ -3,7 +3,8 @@
 from .mean import Mean
 from .constant_mean import ConstantMean
 from .constant_mean_grad import ConstantMeanGrad
+from .linear_mean import LinearMean
 from .multitask_mean import MultitaskMean
 from .zero_mean import ZeroMean
 
-__all__ = ["Mean", "ConstantMean", "ConstantMeanGrad", "MultitaskMean", "ZeroMean"]
+__all__ = ["Mean", "ConstantMean", "ConstantMeanGrad", "LinearMean", "MultitaskMean", "ZeroMean"]
diff --git a/gpytorch/means/linear_mean.py b/gpytorch/means/linear_mean.py
new file mode 100644
index 00000000..d6cf47e7
--- /dev/null
+++ b/gpytorch/means/linear_mean.py
@@ -0,0 +1,19 @@
+#!/usr/bin/env python3
+
+import torch
+from .mean import Mean
+
+
+class LinearMean(Mean):
+    def __init__(self, input_size, batch_shape=torch.Size(), bias=True):
+        super().__init__()
+        self.register_parameter(name='weights',
+                                parameter=torch.nn.Parameter(torch.randn(*batch_shape, input_size, 1)))
+        if bias:
+            self.register_parameter(name='bias', parameter=torch.nn.Parameter(torch.randn(*batch_shape, 1)))
+
+    def forward(self, x):
+        res = x.matmul(self.weights).squeeze(-1)
+        if self.bias is not None:
+            res = res + self.bias
+        return res
diff --git a/gpytorch/models/exact_gp.py b/gpytorch/models/exact_gp.py
index 65c4244d..5f2679ce 100644
--- a/gpytorch/models/exact_gp.py
+++ b/gpytorch/models/exact_gp.py
@@ -91,28 +91,24 @@ class ExactGP(GP):
 
         .. note::
             If `targets` is a batch (e.g. `b x m`), then the GP returned from this method will be a batch mode GP.
-            If `inputs` is of the same (or lesser) dimension as `targets`, then it is assumed that the fantasy points
-            are the same for each target batch.
 
         Args:
-            - :attr:`inputs` (Tensor `b1 x ... x bk x m x d` or `f x b1 x ... x bk x m x d`): Locations of fantasy
-                observations.
-            - :attr:`targets` (Tensor `b1 x ... x bk x m` or `f x b1 x ... x bk x m`): Labels of fantasy observations.
-
+            - :attr:`inputs` (Tensor `m x d` or `b x m x d`): Locations of fantasy observations.
+            - :attr:`targets` (Tensor `m` or `b x m`): Labels of fantasy observations.
         Returns:
             - :class:`ExactGP`
                 An `ExactGP` model with `n + m` training examples, where the `m` fantasy examples have been added
                 and all test-time caches have been updated.
         """
         if self.prediction_strategy is None:
-            raise RuntimeError(
-                "Fantasy observations can only be added after making predictions with a model so that "
-                "all test independent caches exist. Call the model on some data first!"
-            )
+            raise RuntimeError("Fantasy observations can only be added after making predictions with a model so that "
+                               "all test independent caches exist. Call the model on some data first!")
 
-        model_batch_shape = self.train_inputs[0].shape[:-2]
+        if self.train_inputs[0].dim() > 2:
+            raise RuntimeError("Adding fantasy observations to a GP that is already batch mode will not be supported "
+                               "until GPyTorch supports multiple batch dimensions.")
 
-        if self.train_targets.dim() > len(model_batch_shape) + 1:
+        if self.train_targets.dim() > 1:
             raise RuntimeError("Cannot yet add fantasy observations to multitask GPs, but this is coming soon!")
 
         if not isinstance(inputs, list):
@@ -120,39 +116,22 @@ class ExactGP(GP):
 
         inputs = [i.unsqueeze(-1) if i.ndimension() == 1 else i for i in inputs]
 
-        target_batch_shape = targets.shape[:-1]
-        input_batch_shape = inputs[0].shape[:-2]
-        tbdim, ibdim = len(target_batch_shape), len(input_batch_shape)
-
-        if not (tbdim == ibdim + 1 or tbdim == ibdim):
-            raise RuntimeError(
-                f"Unsupported batch shapes: The target batch shape ({target_batch_shape}) must have either the "
-                f"same dimension as or one more dimension than the input batch shape ({input_batch_shape})"
-            )
-
-        # Check whether we can properly broadcast batch dimensions
-        err_msg = (
-            f"Model batch shape ({model_batch_shape}) and target batch shape "
-            f"({target_batch_shape}) are not broadcastable."
-        )
-        _mul_broadcast_shape(model_batch_shape, target_batch_shape, error_msg=err_msg)
-
-        if len(model_batch_shape) > len(input_batch_shape):
-            input_batch_shape = model_batch_shape
-        if len(model_batch_shape) > len(target_batch_shape):
-            target_batch_shape = model_batch_shape
+        # If input is n x d but targets is b x n x d, expand input to b x n x d
+        for i, input in enumerate(inputs):
+            if input.dim() == targets.dim():
+                inputs[i] = input.unsqueeze(0).repeat(*targets.shape[:-1], 1, 1)
+                batch_shape = inputs[i].shape[:-2]
+            elif input.dim() == targets.dim() + 1:
+                batch_shape = input.shape[:-2]
+            elif input.dim() < targets.dim():
+                batch_shape = targets.shape[:-2]
+                inputs[i] = input.expand(*batch_shape, *input.shape[-2:])
 
-        # If input has no fantasy batch dimension but target does, we can save memory and computation by not
-        # computing the covariance for each element of the batch. Therefore we don't expand the inputs to the
-        # size of the fantasy model here - this is done below, after the evaluation and fast fantasy update
-        train_inputs = [tin.expand(input_batch_shape + tin.shape[-2:]) for tin in self.train_inputs]
-        train_targets = self.train_targets.expand(target_batch_shape + self.train_targets.shape[-1:])
+        train_inputs = [tin.expand(*batch_shape, *tin.shape[-2:]) for tin in list(self.train_inputs)]
+        train_targets = self.train_targets.expand(*batch_shape, *self.train_targets.shape[-2:])
 
-        full_inputs = [
-            torch.cat([train_input, input.expand(input_batch_shape + input.shape[-2:])], dim=-2)
-            for train_input, input in zip(train_inputs, inputs)
-        ]
-        full_targets = torch.cat([train_targets, targets.expand(target_batch_shape + targets.shape[-1:])], dim=-1)
+        full_inputs = [torch.cat([train_input, input], dim=-2) for train_input, input in zip(train_inputs, inputs)]
+        full_targets = torch.cat([train_targets, targets], dim=-1)
 
         try:
             fantasy_kwargs = {"noise": kwargs.pop("noise")}
@@ -176,6 +155,8 @@ class ExactGP(GP):
         self.train_targets = old_train_targets
         self.likelihood = old_likelihood
 
+        new_model.train_inputs = full_inputs
+        new_model.train_targets = full_targets
         new_model.likelihood = old_likelihood.get_fantasy_likelihood(**fantasy_kwargs)
         new_model.prediction_strategy = old_pred_strat.get_fantasy_strategy(
             inputs,
@@ -186,13 +167,6 @@ class ExactGP(GP):
             **fantasy_kwargs,
         )
 
-        # if the fantasies are at the same points, we need to expand the inputs for the new model
-        if tbdim == ibdim + 1:
-            new_model.train_inputs = [fi.expand(target_batch_shape + fi.shape[-2:]) for fi in full_inputs]
-        else:
-            new_model.train_inputs = full_inputs
-        new_model.train_targets = full_targets
-
         return new_model
 
     def train(self, mode=True):
@@ -262,7 +236,7 @@ class ExactGP(GP):
             # Concatenate the input to the training input
             full_inputs = []
             batch_shape = train_inputs[0].shape[:-2]
-            for train_input, input in zip(train_inputs, inputs):
+            for i, (train_input, input) in enumerate(zip(train_inputs, inputs)):
                 # Make sure the batch shapes agree for training/test data
                 if batch_shape != train_input.shape[:-2]:
                     batch_shape = _mul_broadcast_shape(batch_shape, train_input.shape[:-2])
diff --git a/gpytorch/models/exact_prediction_strategies.py b/gpytorch/models/exact_prediction_strategies.py
index df410df0..ff2b6976 100644
--- a/gpytorch/models/exact_prediction_strategies.py
+++ b/gpytorch/models/exact_prediction_strategies.py
@@ -5,7 +5,7 @@ import torch
 from .. import settings
 from ..lazy import (
     InterpolatedLazyTensor, MatmulLazyTensor, RootLazyTensor, SumLazyTensor, ZeroLazyTensor,
-    lazify, delazify, LazyEvaluatedKernelTensor, NonLazyTensor, BatchRepeatLazyTensor
+    lazify, delazify, LazyEvaluatedKernelTensor
 )
 from ..utils.interpolation import left_interp, left_t_interp
 from ..utils.memoize import cached, add_to_cache
@@ -94,15 +94,11 @@ class DefaultPredictionStrategy(object):
         GP model, use the :meth:`~gpytorch.models.ExactGP.get_fantasy_model` method.
 
         Args:
-            - :attr:`inputs` (Tensor `b1 x ... x bk x m x d` or `f x b1 x ... x bk x m x d`): Locations of fantasy
-                observations.
-            - :attr:`targets` (Tensor `b1 x ... x bk x m` or `f x b1 x ... x bk x m`): Labels of fantasy observations.
-            - :attr:`full_inputs` (Tensor `b1 x ... x bk x n+m x d` or `f x b1 x ... x bk x n+m x d`): Training data
-                concatenated with fantasy inputs
-            - :attr:`full_targets` (Tensor `b1 x ... x bk x n+m` or `f x b1 x ... x bk x n+m`): Training labels
-                concatenated with fantasy labels.
+            - :attr:`inputs` (Tensor `m x d` or `b x m x d`): Locations of fantasy observations.
+            - :attr:`targets` (Tensor `m` or `b x m`): Labels of fantasy observations.
+            - :attr:`full_inputs` (Tensor `n+m x d` or `b x n+m x d`): Training data concatenated with fantasy inputs
+            - :attr:`full_targets` (Tensor `n+m` or `b x n+m`): Training labels concatenated with fantasy labels.
             - :attr:`full_output` (:class:`gpytorch.distributions.MultivariateNormal`): Prior called on full_inputs
-
         Returns:
             - :class:`DefaultPredictionStrategy`
                 A `DefaultPredictionStrategy` model with `n + m` training examples, where the `m` fantasy examples have
@@ -145,18 +141,12 @@ class DefaultPredictionStrategy(object):
 
         # Solve for "b", the lower portion of the *new* \\alpha corresponding to the fantasy points.
         schur_complement = fant_fant_covar - fant_train_covar.matmul(fant_solve)
-        if self.mean_cache.dim() == 1:
-            ftcm = fant_train_covar.matmul(self.mean_cache)
-        elif full_inputs[0].dim() <= full_targets.dim():
-            ftcm = torch.einsum("...ij,...j->...i", [fant_train_covar, self.mean_cache])
-        else:
-            ftcm = torch.einsum("f...ij,...j->f...i", [fant_train_covar, self.mean_cache])
-        small_system_rhs = targets - fant_mean - ftcm
+        small_system_rhs = targets - fant_mean - fant_train_covar.matmul(self.mean_cache)
         small_system_rhs = small_system_rhs.unsqueeze(-1)
         # Schur complement of a spd matrix is guaranteed to be positive definite
         if small_system_rhs.requires_grad or schur_complement.requires_grad:
             # TODO: Delete this part of the if statement when PyTorch implements cholesky_solve derivative.
-            fant_cache_lower = torch.solve(small_system_rhs, schur_complement)[0]
+            fant_cache_lower = torch.gesv(small_system_rhs, schur_complement)[0]
         else:
             fant_cache_lower = cholesky_solve(small_system_rhs, torch.cholesky(schur_complement))
 
@@ -211,22 +201,11 @@ class DefaultPredictionStrategy(object):
         cap_mat = new_root.transpose(-2, -1).matmul(new_root)
         if cap_mat.requires_grad or new_root.requires_grad:
             # TODO: Delete this part of the if statement when PyTorch implements cholesky_solve derivative.
-            new_covar_cache = torch.solve(new_root.transpose(-2, -1), cap_mat)[0].transpose(-2, -1)
+            new_covar_cache = torch.gesv(new_root.transpose(-2, -1), cap_mat)[0].transpose(-2, -1)
         else:
             new_covar_cache = cholesky_solve(new_root.transpose(-2, -1), torch.cholesky(cap_mat))
             new_covar_cache = new_covar_cache.transpose(-2, -1)
 
-        # Expand inputs accordingly if necessary (for fantasies at the same points)
-        if full_inputs[0].dim() <= full_targets.dim():
-            fant_batch_shape = full_targets.shape[:1]
-            n_batch = len(full_mean.shape[:-1])
-            repeat_shape = fant_batch_shape + torch.Size([1] * n_batch)
-            full_inputs = [fi.expand(fant_batch_shape + fi.shape) for fi in full_inputs]
-            full_mean = full_mean.expand(fant_batch_shape + full_mean.shape)
-            full_covar = BatchRepeatLazyTensor(full_covar, repeat_shape)
-            new_root = BatchRepeatLazyTensor(NonLazyTensor(new_root), repeat_shape)
-            # no need to repeat the covar cache, broadcasting will do the right thing
-
         # Create new DefaultPredictionStrategy object
         fant_strat = self.__class__(
             train_inputs=full_inputs,
@@ -236,7 +215,7 @@ class DefaultPredictionStrategy(object):
             root=new_root,
             inv_root=new_covar_cache,
         )
-        fant_strat._memoize_cache = {"mean_cache": fant_mean_cache, "covar_cache": new_covar_cache}
+        setattr(fant_strat, "_memoize_cache", {"mean_cache": fant_mean_cache, "covar_cache": new_covar_cache})
 
         return fant_strat
 
diff --git a/gpytorch/utils/linear_cg.py b/gpytorch/utils/linear_cg.py
index 35af8f7d..a6d8393a 100644
--- a/gpytorch/utils/linear_cg.py
+++ b/gpytorch/utils/linear_cg.py
@@ -284,7 +284,7 @@ def linear_cg(
             prev_beta.copy_(beta_tridiag)
 
     # Un-normalize
-    result = result.mul(rhs_norm)
+    result.mul_(rhs_norm)
 
     if not tolerance_reached and n_iter > 0:
         warnings.warn(


Test patch:
diff --git a/test/means/test_linear_mean.py b/test/means/test_linear_mean.py
new file mode 100644
index 00000000..c1b033b7
--- /dev/null
+++ b/test/means/test_linear_mean.py
@@ -0,0 +1,36 @@
+#!/usr/bin/env python3
+
+import torch
+import unittest
+from test.means._base_mean_test_case import BaseMeanTestCase
+from gpytorch.means import LinearMean
+
+
+class TestLinearMean(BaseMeanTestCase, unittest.TestCase):
+
+    def create_mean(self, input_size=1, batch_shape=torch.Size(), bias=True, **kwargs):
+        return LinearMean(input_size=input_size, batch_shape=batch_shape, bias=bias)
+
+    def forward_vec(self):
+        n = 4
+        test_x = torch.randn(n)
+        mean = self.create_mean(input_size=1)
+        self.assertEqual(mean(test_x).shape, torch.Size([4]))
+
+    def test_forward_mat(self):
+        n, d = 4, 5
+        test_x = torch.randn(n, d)
+        mean = self.create_mean(d)
+        self.assertEqual(mean(test_x).shape, torch.Size([n]))
+
+    def test_forward_mat_batch(self):
+        b, n, d = torch.Size([3]), 4, 5
+        test_x = torch.randn(*b, n, d)
+        mean = self.create_mean(d, b)
+        self.assertEqual(mean(test_x).shape, torch.Size([*b, n]))
+
+    def test_forward_mat_multi_batch(self):
+        b, n, d = torch.Size([2, 3]), 4, 5
+        test_x = torch.randn(*b, n, d)
+        mean = self.create_mean(d, b)
+        self.assertEqual(mean(test_x).shape, torch.Size([*b, n]))