We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[Bug] size mismatch when using natural gradient in multi-output GP
# 🐛 Bug

size mismatch when using natural gradient in multi-output GP

## To reproduce

** Code snippet to reproduce **
```python
import math
import torch
import gpytorch
import tqdm
from matplotlib import pyplot as plt

train_x = torch.linspace(0, 1, 100)

train_y = torch.stack([
    torch.sin(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * 0.2,
    torch.cos(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * 0.2,
    torch.sin(train_x * (2 * math.pi)) + 2 * torch.cos(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * 0.2,
    -torch.cos(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * 0.2,
], -1)

print(train_x.shape, train_y.shape)

num_latents = 3
num_tasks = 4


class IndependentMultitaskGPModel(gpytorch.models.ApproximateGP):
    def __init__(self):
        # Let's use a different set of inducing points for each task
        inducing_points = torch.rand(num_tasks, 16, 1)

        # We have to mark the CholeskyVariationalDistribution as batch
        # so that we learn a variational distribution for each task
        # variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(
        #     inducing_points.size(-2), batch_shape=torch.Size([num_tasks])
        # )
        variational_distribution = gpytorch.variational.NaturalVariationalDistribution(
            inducing_points.size(-2), batch_shape=torch.Size([num_tasks])
        )

        variational_strategy = gpytorch.variational.IndependentMultitaskVariationalStrategy(
            gpytorch.variational.VariationalStrategy(
                self, inducing_points, variational_distribution, learn_inducing_locations=True
            ),
            num_tasks=4,
        )

        super().__init__(variational_strategy)

        # The mean and covariance modules should be marked as batch
        # so we learn a different set of hyperparameters
        self.mean_module = gpytorch.means.ConstantMean(batch_shape=torch.Size([num_tasks]))
        self.covar_module = gpytorch.kernels.ScaleKernel(
            gpytorch.kernels.RBFKernel(batch_shape=torch.Size([num_tasks])),
            batch_shape=torch.Size([num_tasks])
        )

    def forward(self, x):
        # The forward function should be written as if we were dealing with each output
        # dimension in batch
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)


# model = MultitaskGPModel()

model = IndependentMultitaskGPModel()
likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=num_tasks)

# this is for running the notebook in our testing framework
import os
smoke_test = ('CI' in os.environ)
num_epochs = 1 if smoke_test else 500


model.train()
likelihood.train()

ngd_optimizer = gpytorch.optim.NGD(model.variational_parameters(),
                                   num_data=train_y.size(0),
                                   lr=0.1)
# We use SGD here, rather than Adam. Emperically, we find that SGD is better for variational regression
optimizer = torch.optim.Adam([
    {'params': model.hyperparameters()},
    {'params': likelihood.parameters()},
], lr=0.1)

# Our loss object. We're using the VariationalELBO, which essentially just computes the ELBO
mll = gpytorch.mlls.VariationalELBO(likelihood, model, num_data=train_y.size(0))

# We use more CG iterations here because the preconditioner introduced in the NeurIPS paper seems to be less
# effective for VI.
epochs_iter = tqdm.tqdm(range(num_epochs), desc="Epoch")
for i in epochs_iter:
    # Within each iteration, we will go over each minibatch of data
    ngd_optimizer.zero_grad()
    optimizer.zero_grad()
    output = model(train_x)
    loss = -mll(output, train_y)
    epochs_iter.set_postfix(loss=loss.item())
    loss.backward()
    ngd_optimizer.step()
    optimizer.step()

# Set into eval mode
model.eval()
likelihood.eval()

# Initialize plots
fig, axs = plt.subplots(1, num_tasks, figsize=(4 * num_tasks, 3))

# Make predictions
with torch.no_grad(), gpytorch.settings.fast_pred_var():
    test_x = torch.linspace(0, 1, 51)
    predictions = likelihood(model(test_x))
    mean = predictions.mean
    lower, upper = predictions.confidence_region()

for task, ax in enumerate(axs):
    # Plot training data as black stars
    ax.plot(train_x.detach().numpy(), train_y[:, task].detach().numpy(), 'k*')
    # Predictive mean as blue line
    ax.plot(test_x.numpy(), mean[:, task].numpy(), 'b')
    # Shade in confidence
    ax.fill_between(test_x.numpy(), lower[:, task].numpy(), upper[:, task].numpy(), alpha=0.5)
    ax.set_ylim([-3, 3])
    ax.legend(['Observed Data', 'Mean', 'Confidence'])
    ax.set_title(f'Task {task + 1}')

fig.tight_layout()
plt.show()
```

** Stack trace/error message **
```bash
Traceback (most recent call last):
torch.Size([100]) torch.Size([100, 4])
Epoch:   0%|                                            | 0/500 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "multi-svgp.py", line 99, in <module>
    output = model(train_x)
  File "/home/user/miniconda2/envs/torch_play/lib/python3.7/site-packages/gpytorch/models/approximate_gp.py", line 81, in __call__
    return self.variational_strategy(inputs, prior=prior)
  File "/home/user/miniconda2/envs/torch_play/lib/python3.7/site-packages/gpytorch/variational/independent_multitask_variational_strategy.py", line 47, in __call__
    function_dist = self.base_variational_strategy(x, prior=prior)
  File "/home/user/miniconda2/envs/torch_play/lib/python3.7/site-packages/gpytorch/variational/variational_strategy.py", line 166, in __call__
    return super().__call__(x, prior=prior)
  File "/home/user/miniconda2/envs/torch_play/lib/python3.7/site-packages/gpytorch/variational/_variational_strategy.py", line 114, in __call__
    self._variational_distribution.initialize_variational_distribution(prior_dist)
  File "/home/user/miniconda2/envs/torch_play/lib/python3.7/site-packages/gpytorch/variational/natural_variational_distribution.py", line 69, in initialize_variational_distribution
    self.natural_vec.data.copy_((prior_prec @ prior_mean).add_(noise))
RuntimeError: size mismatch, m1: [64 x 16], m2: [4 x 16] at /pytorch/aten/src/TH/generic/THTensorMath.cpp:41
```

## Expected Behavior

Getting similar results as in the [non-natural gradient notebook](https://docs.gpytorch.ai/en/v1.2.0/examples/04_Variational_and_Approximate_GPs/SVGP_Multitask_GP_Regression.html#Train-the-model)

## System information

**Please complete the following information:**
-  GPyTorch Version 1.2.0 
- PyTorch Version 1.6.0
-  Debian 10

## Additional context
I followed the [NGD tutorial](https://docs.gpytorch.ai/en/v1.2.0/examples/04_Variational_and_Approximate_GPs/Natural_Gradient_Descent.html) and replace the necessary code in the multi-output SVGP cases. Seems the problems come to the shape of `prior_prec` and `prior_mean`. 


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `gpytorch/variational/natural_variational_distribution.py`
Content:
```
1 #!/usr/bin/env python3
2 
3 import abc
4 
5 import torch
6 
7 from ..distributions import MultivariateNormal
8 from ..lazy import CholLazyTensor, TriangularLazyTensor
9 from ._variational_distribution import _VariationalDistribution
10 
11 
12 class _NaturalVariationalDistribution(_VariationalDistribution, abc.ABC):
13     r"""Any :obj:`~gpytorch.variational._VariationalDistribution` which calculates
14     natural gradients with respect to its parameters.
15     """
16     pass
17 
18 
19 class NaturalVariationalDistribution(_NaturalVariationalDistribution):
20     r"""A multivariate normal :obj:`~gpytorch.variational._VariationalDistribution`,
21     parameterized by **natural** parameters.
22 
23     .. note::
24        The :obj:`~gpytorch.variational.NaturalVariationalDistribution` can only
25        be used with :obj:`gpytorch.optim.NGD`, or other optimizers
26        that follow exactly the gradient direction. Failure to do so will cause
27        the natural matrix :math:`\mathbf \Theta_\text{mat}` to stop being
28        positive definite, and a :obj:`~RuntimeError` will be raised.
29 
30     .. seealso::
31         The `natural gradient descent tutorial
32         <examples/04_Variational_and_Approximate_GPs/Natural_Gradient_Descent.ipynb>`_
33         for use instructions.
34 
35         The :obj:`~gpytorch.variational.TrilNaturalVariationalDistribution` for
36         a more numerically stable parameterization, at the cost of needing more
37         iterations to make variational regression converge.
38 
39     :param int num_inducing_points: Size of the variational distribution. This implies that the variational mean
40         should be this size, and the variational covariance matrix should have this many rows and columns.
41     :param batch_shape: Specifies an optional batch size
42         for the variational parameters. This is useful for example when doing additive variational inference.
43     :type batch_shape: :obj:`torch.Size`, optional
44     :param float mean_init_std: (Default: 1e-3) Standard deviation of gaussian noise to add to the mean initialization.
45 
46     """
47 
48     def __init__(self, num_inducing_points, batch_shape=torch.Size([]), mean_init_std=1e-3, **kwargs):
49         super().__init__(num_inducing_points=num_inducing_points, batch_shape=batch_shape, mean_init_std=mean_init_std)
50         scaled_mean_init = torch.zeros(num_inducing_points)
51         neg_prec_init = torch.eye(num_inducing_points, num_inducing_points).mul(-0.5)
52         scaled_mean_init = scaled_mean_init.repeat(*batch_shape, 1)
53         neg_prec_init = neg_prec_init.repeat(*batch_shape, 1, 1)
54 
55         # eta1 and eta2 parameterization of the variational distribution
56         self.register_parameter(name="natural_vec", parameter=torch.nn.Parameter(scaled_mean_init))
57         self.register_parameter(name="natural_mat", parameter=torch.nn.Parameter(neg_prec_init))
58 
59     def forward(self):
60         mean, chol_covar = _NaturalToMuVarSqrt.apply(self.natural_vec, self.natural_mat)
61         res = MultivariateNormal(mean, CholLazyTensor(TriangularLazyTensor(chol_covar)))
62         return res
63 
64     def initialize_variational_distribution(self, prior_dist):
65         prior_prec = prior_dist.covariance_matrix.inverse()
66         prior_mean = prior_dist.mean
67         noise = torch.randn_like(prior_mean).mul_(self.mean_init_std)
68 
69         self.natural_vec.data.copy_((prior_prec @ prior_mean).add_(noise))
70         self.natural_mat.data.copy_(prior_prec.mul(-0.5))
71 
72 
73 def _triangular_inverse(A, upper=False):
74     eye = torch.eye(A.size(-1), dtype=A.dtype, device=A.device)
75     return eye.triangular_solve(A, upper=upper).solution
76 
77 
78 def _phi_for_cholesky_(A):
79     "Modifies A to be the phi function used in differentiating through Cholesky"
80     A.tril_().diagonal(offset=0, dim1=-2, dim2=-1).mul_(0.5)
81     return A
82 
83 
84 def _cholesky_backward(dout_dL, L, L_inverse):
85     # c.f. https://github.com/pytorch/pytorch/blob/25ba802ce4cbdeaebcad4a03cec8502f0de9b7b3/
86     #      tools/autograd/templates/Functions.cpp
87     A = L.transpose(-1, -2) @ dout_dL
88     phi = _phi_for_cholesky_(A)
89     grad_input = (L_inverse.transpose(-1, -2) @ phi) @ L_inverse
90     # Symmetrize gradient
91     return grad_input.add(grad_input.transpose(-1, -2)).mul_(0.5)
92 
93 
94 class _NaturalToMuVarSqrt(torch.autograd.Function):
95     @staticmethod
96     def _forward(nat_mean, nat_covar):
97         try:
98             L_inv = torch.cholesky(-2.0 * nat_covar, upper=False)
99         except RuntimeError as e:
100             if str(e).startswith("cholesky"):
101                 raise RuntimeError(
102                     "Non-negative-definite natural covariance. You probably "
103                     "updated it using an optimizer other than gpytorch.optim.NGD (such as Adam). "
104                     "This is not supported."
105                 )
106             else:
107                 raise e
108         L = _triangular_inverse(L_inv, upper=False)
109         S = L.transpose(-1, -2) @ L
110         mu = (S @ nat_mean.unsqueeze(-1)).squeeze(-1)
111         # Two choleskys are annoying, but we don't have good support for a
112         # LazyTensor of form L.T @ L
113         return mu, torch.cholesky(S, upper=False)
114 
115     @staticmethod
116     def forward(ctx, nat_mean, nat_covar):
117         mu, L = _NaturalToMuVarSqrt._forward(nat_mean, nat_covar)
118         ctx.save_for_backward(mu, L)
119         return mu, L
120 
121     @staticmethod
122     def _backward(dout_dmu, dout_dL, mu, L, C):
123         """Calculate dout/d(eta1, eta2), which are:
124         eta1 = mu
125         eta2 = mu*mu^T + LL^T = mu*mu^T + Sigma
126 
127         Thus:
128         dout/deta1 = dout/dmu + dout/dL dL/deta1
129         dout/deta2 = dout/dL dL/deta1
130 
131         For L = chol(eta2 - eta1*eta1^T).
132         dout/dSigma = _cholesky_backward(dout/dL, L)
133         dout/deta2 = dout/dSigma
134         dSigma/deta1 = -2* (dout/dSigma) mu
135         """
136         dout_dSigma = _cholesky_backward(dout_dL, L, C)
137         dout_deta1 = dout_dmu - 2 * (dout_dSigma @ mu.unsqueeze(-1)).squeeze(-1)
138         return dout_deta1, dout_dSigma
139 
140     @staticmethod
141     def backward(ctx, dout_dmu, dout_dL):
142         "Calculates the natural gradient with respect to nat_mean, nat_covar"
143         mu, L = ctx.saved_tensors
144         C = _triangular_inverse(L, upper=False)
145         return _NaturalToMuVarSqrt._backward(dout_dmu, dout_dL, mu, L, C)
146
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/gpytorch/variational/natural_variational_distribution.py b/gpytorch/variational/natural_variational_distribution.py
--- a/gpytorch/variational/natural_variational_distribution.py
+++ b/gpytorch/variational/natural_variational_distribution.py
@@ -6,6 +6,7 @@
 
 from ..distributions import MultivariateNormal
 from ..lazy import CholLazyTensor, TriangularLazyTensor
+from ..utils.cholesky import psd_safe_cholesky
 from ._variational_distribution import _VariationalDistribution
 
 
@@ -66,7 +67,7 @@
         prior_mean = prior_dist.mean
         noise = torch.randn_like(prior_mean).mul_(self.mean_init_std)
 
-        self.natural_vec.data.copy_((prior_prec @ prior_mean).add_(noise))
+        self.natural_vec.data.copy_((prior_prec @ prior_mean.unsqueeze(-1)).squeeze(-1).add_(noise))
         self.natural_mat.data.copy_(prior_prec.mul(-0.5))
 
 
@@ -95,7 +96,7 @@
     @staticmethod
     def _forward(nat_mean, nat_covar):
         try:
-            L_inv = torch.cholesky(-2.0 * nat_covar, upper=False)
+            L_inv = psd_safe_cholesky(-2.0 * nat_covar, upper=False)
         except RuntimeError as e:
             if str(e).startswith("cholesky"):
                 raise RuntimeError(
@@ -110,7 +111,7 @@
         mu = (S @ nat_mean.unsqueeze(-1)).squeeze(-1)
         # Two choleskys are annoying, but we don't have good support for a
         # LazyTensor of form L.T @ L
-        return mu, torch.cholesky(S, upper=False)
+        return mu, psd_safe_cholesky(S, upper=False)
 
     @staticmethod
     def forward(ctx, nat_mean, nat_covar):

{"golden_diff": "diff --git a/gpytorch/variational/natural_variational_distribution.py b/gpytorch/variational/natural_variational_distribution.py\n--- a/gpytorch/variational/natural_variational_distribution.py\n+++ b/gpytorch/variational/natural_variational_distribution.py\n@@ -6,6 +6,7 @@\n \n from ..distributions import MultivariateNormal\n from ..lazy import CholLazyTensor, TriangularLazyTensor\n+from ..utils.cholesky import psd_safe_cholesky\n from ._variational_distribution import _VariationalDistribution\n \n \n@@ -66,7 +67,7 @@\n         prior_mean = prior_dist.mean\n         noise = torch.randn_like(prior_mean).mul_(self.mean_init_std)\n \n-        self.natural_vec.data.copy_((prior_prec @ prior_mean).add_(noise))\n+        self.natural_vec.data.copy_((prior_prec @ prior_mean.unsqueeze(-1)).squeeze(-1).add_(noise))\n         self.natural_mat.data.copy_(prior_prec.mul(-0.5))\n \n \n@@ -95,7 +96,7 @@\n     @staticmethod\n     def _forward(nat_mean, nat_covar):\n         try:\n-            L_inv = torch.cholesky(-2.0 * nat_covar, upper=False)\n+            L_inv = psd_safe_cholesky(-2.0 * nat_covar, upper=False)\n         except RuntimeError as e:\n             if str(e).startswith(\"cholesky\"):\n                 raise RuntimeError(\n@@ -110,7 +111,7 @@\n         mu = (S @ nat_mean.unsqueeze(-1)).squeeze(-1)\n         # Two choleskys are annoying, but we don't have good support for a\n         # LazyTensor of form L.T @ L\n-        return mu, torch.cholesky(S, upper=False)\n+        return mu, psd_safe_cholesky(S, upper=False)\n \n     @staticmethod\n     def forward(ctx, nat_mean, nat_covar):\n", "issue": "[Bug] size mismatch when using natural gradient in multi-output GP\n# \ud83d\udc1b Bug\r\n\r\nsize mismatch when using natural gradient in multi-output GP\r\n\r\n## To reproduce\r\n\r\n** Code snippet to reproduce **\r\n```python\r\nimport math\r\nimport torch\r\nimport gpytorch\r\nimport tqdm\r\nfrom matplotlib import pyplot as plt\r\n\r\ntrain_x = torch.linspace(0, 1, 100)\r\n\r\ntrain_y = torch.stack([\r\n    torch.sin(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * 0.2,\r\n    torch.cos(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * 0.2,\r\n    torch.sin(train_x * (2 * math.pi)) + 2 * torch.cos(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * 0.2,\r\n    -torch.cos(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * 0.2,\r\n], -1)\r\n\r\nprint(train_x.shape, train_y.shape)\r\n\r\nnum_latents = 3\r\nnum_tasks = 4\r\n\r\n\r\nclass IndependentMultitaskGPModel(gpytorch.models.ApproximateGP):\r\n    def __init__(self):\r\n        # Let's use a different set of inducing points for each task\r\n        inducing_points = torch.rand(num_tasks, 16, 1)\r\n\r\n        # We have to mark the CholeskyVariationalDistribution as batch\r\n        # so that we learn a variational distribution for each task\r\n        # variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(\r\n        #     inducing_points.size(-2), batch_shape=torch.Size([num_tasks])\r\n        # )\r\n        variational_distribution = gpytorch.variational.NaturalVariationalDistribution(\r\n            inducing_points.size(-2), batch_shape=torch.Size([num_tasks])\r\n        )\r\n\r\n        variational_strategy = gpytorch.variational.IndependentMultitaskVariationalStrategy(\r\n            gpytorch.variational.VariationalStrategy(\r\n                self, inducing_points, variational_distribution, learn_inducing_locations=True\r\n            ),\r\n            num_tasks=4,\r\n        )\r\n\r\n        super().__init__(variational_strategy)\r\n\r\n        # The mean and covariance modules should be marked as batch\r\n        # so we learn a different set of hyperparameters\r\n        self.mean_module = gpytorch.means.ConstantMean(batch_shape=torch.Size([num_tasks]))\r\n        self.covar_module = gpytorch.kernels.ScaleKernel(\r\n            gpytorch.kernels.RBFKernel(batch_shape=torch.Size([num_tasks])),\r\n            batch_shape=torch.Size([num_tasks])\r\n        )\r\n\r\n    def forward(self, x):\r\n        # The forward function should be written as if we were dealing with each output\r\n        # dimension in batch\r\n        mean_x = self.mean_module(x)\r\n        covar_x = self.covar_module(x)\r\n        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)\r\n\r\n\r\n# model = MultitaskGPModel()\r\n\r\nmodel = IndependentMultitaskGPModel()\r\nlikelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=num_tasks)\r\n\r\n# this is for running the notebook in our testing framework\r\nimport os\r\nsmoke_test = ('CI' in os.environ)\r\nnum_epochs = 1 if smoke_test else 500\r\n\r\n\r\nmodel.train()\r\nlikelihood.train()\r\n\r\nngd_optimizer = gpytorch.optim.NGD(model.variational_parameters(),\r\n                                   num_data=train_y.size(0),\r\n                                   lr=0.1)\r\n# We use SGD here, rather than Adam. Emperically, we find that SGD is better for variational regression\r\noptimizer = torch.optim.Adam([\r\n    {'params': model.hyperparameters()},\r\n    {'params': likelihood.parameters()},\r\n], lr=0.1)\r\n\r\n# Our loss object. We're using the VariationalELBO, which essentially just computes the ELBO\r\nmll = gpytorch.mlls.VariationalELBO(likelihood, model, num_data=train_y.size(0))\r\n\r\n# We use more CG iterations here because the preconditioner introduced in the NeurIPS paper seems to be less\r\n# effective for VI.\r\nepochs_iter = tqdm.tqdm(range(num_epochs), desc=\"Epoch\")\r\nfor i in epochs_iter:\r\n    # Within each iteration, we will go over each minibatch of data\r\n    ngd_optimizer.zero_grad()\r\n    optimizer.zero_grad()\r\n    output = model(train_x)\r\n    loss = -mll(output, train_y)\r\n    epochs_iter.set_postfix(loss=loss.item())\r\n    loss.backward()\r\n    ngd_optimizer.step()\r\n    optimizer.step()\r\n\r\n# Set into eval mode\r\nmodel.eval()\r\nlikelihood.eval()\r\n\r\n# Initialize plots\r\nfig, axs = plt.subplots(1, num_tasks, figsize=(4 * num_tasks, 3))\r\n\r\n# Make predictions\r\nwith torch.no_grad(), gpytorch.settings.fast_pred_var():\r\n    test_x = torch.linspace(0, 1, 51)\r\n    predictions = likelihood(model(test_x))\r\n    mean = predictions.mean\r\n    lower, upper = predictions.confidence_region()\r\n\r\nfor task, ax in enumerate(axs):\r\n    # Plot training data as black stars\r\n    ax.plot(train_x.detach().numpy(), train_y[:, task].detach().numpy(), 'k*')\r\n    # Predictive mean as blue line\r\n    ax.plot(test_x.numpy(), mean[:, task].numpy(), 'b')\r\n    # Shade in confidence\r\n    ax.fill_between(test_x.numpy(), lower[:, task].numpy(), upper[:, task].numpy(), alpha=0.5)\r\n    ax.set_ylim([-3, 3])\r\n    ax.legend(['Observed Data', 'Mean', 'Confidence'])\r\n    ax.set_title(f'Task {task + 1}')\r\n\r\nfig.tight_layout()\r\nplt.show()\r\n```\r\n\r\n** Stack trace/error message **\r\n```bash\r\nTraceback (most recent call last):\r\ntorch.Size([100]) torch.Size([100, 4])\r\nEpoch:   0%|                                            | 0/500 [00:00<?, ?it/s]\r\nTraceback (most recent call last):\r\n  File \"multi-svgp.py\", line 99, in <module>\r\n    output = model(train_x)\r\n  File \"/home/user/miniconda2/envs/torch_play/lib/python3.7/site-packages/gpytorch/models/approximate_gp.py\", line 81, in __call__\r\n    return self.variational_strategy(inputs, prior=prior)\r\n  File \"/home/user/miniconda2/envs/torch_play/lib/python3.7/site-packages/gpytorch/variational/independent_multitask_variational_strategy.py\", line 47, in __call__\r\n    function_dist = self.base_variational_strategy(x, prior=prior)\r\n  File \"/home/user/miniconda2/envs/torch_play/lib/python3.7/site-packages/gpytorch/variational/variational_strategy.py\", line 166, in __call__\r\n    return super().__call__(x, prior=prior)\r\n  File \"/home/user/miniconda2/envs/torch_play/lib/python3.7/site-packages/gpytorch/variational/_variational_strategy.py\", line 114, in __call__\r\n    self._variational_distribution.initialize_variational_distribution(prior_dist)\r\n  File \"/home/user/miniconda2/envs/torch_play/lib/python3.7/site-packages/gpytorch/variational/natural_variational_distribution.py\", line 69, in initialize_variational_distribution\r\n    self.natural_vec.data.copy_((prior_prec @ prior_mean).add_(noise))\r\nRuntimeError: size mismatch, m1: [64 x 16], m2: [4 x 16] at /pytorch/aten/src/TH/generic/THTensorMath.cpp:41\r\n```\r\n\r\n## Expected Behavior\r\n\r\nGetting similar results as in the [non-natural gradient notebook](https://docs.gpytorch.ai/en/v1.2.0/examples/04_Variational_and_Approximate_GPs/SVGP_Multitask_GP_Regression.html#Train-the-model)\r\n\r\n## System information\r\n\r\n**Please complete the following information:**\r\n-  GPyTorch Version 1.2.0 \r\n- PyTorch Version 1.6.0\r\n-  Debian 10\r\n\r\n## Additional context\r\nI followed the [NGD tutorial](https://docs.gpytorch.ai/en/v1.2.0/examples/04_Variational_and_Approximate_GPs/Natural_Gradient_Descent.html) and replace the necessary code in the multi-output SVGP cases. Seems the problems come to the shape of `prior_prec` and `prior_mean`. \r\n\n", "before_files": [{"content": "#!/usr/bin/env python3\n\nimport abc\n\nimport torch\n\nfrom ..distributions import MultivariateNormal\nfrom ..lazy import CholLazyTensor, TriangularLazyTensor\nfrom ._variational_distribution import _VariationalDistribution\n\n\nclass _NaturalVariationalDistribution(_VariationalDistribution, abc.ABC):\n    r\"\"\"Any :obj:`~gpytorch.variational._VariationalDistribution` which calculates\n    natural gradients with respect to its parameters.\n    \"\"\"\n    pass\n\n\nclass NaturalVariationalDistribution(_NaturalVariationalDistribution):\n    r\"\"\"A multivariate normal :obj:`~gpytorch.variational._VariationalDistribution`,\n    parameterized by **natural** parameters.\n\n    .. note::\n       The :obj:`~gpytorch.variational.NaturalVariationalDistribution` can only\n       be used with :obj:`gpytorch.optim.NGD`, or other optimizers\n       that follow exactly the gradient direction. Failure to do so will cause\n       the natural matrix :math:`\\mathbf \\Theta_\\text{mat}` to stop being\n       positive definite, and a :obj:`~RuntimeError` will be raised.\n\n    .. seealso::\n        The `natural gradient descent tutorial\n        <examples/04_Variational_and_Approximate_GPs/Natural_Gradient_Descent.ipynb>`_\n        for use instructions.\n\n        The :obj:`~gpytorch.variational.TrilNaturalVariationalDistribution` for\n        a more numerically stable parameterization, at the cost of needing more\n        iterations to make variational regression converge.\n\n    :param int num_inducing_points: Size of the variational distribution. This implies that the variational mean\n        should be this size, and the variational covariance matrix should have this many rows and columns.\n    :param batch_shape: Specifies an optional batch size\n        for the variational parameters. This is useful for example when doing additive variational inference.\n    :type batch_shape: :obj:`torch.Size`, optional\n    :param float mean_init_std: (Default: 1e-3) Standard deviation of gaussian noise to add to the mean initialization.\n\n    \"\"\"\n\n    def __init__(self, num_inducing_points, batch_shape=torch.Size([]), mean_init_std=1e-3, **kwargs):\n        super().__init__(num_inducing_points=num_inducing_points, batch_shape=batch_shape, mean_init_std=mean_init_std)\n        scaled_mean_init = torch.zeros(num_inducing_points)\n        neg_prec_init = torch.eye(num_inducing_points, num_inducing_points).mul(-0.5)\n        scaled_mean_init = scaled_mean_init.repeat(*batch_shape, 1)\n        neg_prec_init = neg_prec_init.repeat(*batch_shape, 1, 1)\n\n        # eta1 and eta2 parameterization of the variational distribution\n        self.register_parameter(name=\"natural_vec\", parameter=torch.nn.Parameter(scaled_mean_init))\n        self.register_parameter(name=\"natural_mat\", parameter=torch.nn.Parameter(neg_prec_init))\n\n    def forward(self):\n        mean, chol_covar = _NaturalToMuVarSqrt.apply(self.natural_vec, self.natural_mat)\n        res = MultivariateNormal(mean, CholLazyTensor(TriangularLazyTensor(chol_covar)))\n        return res\n\n    def initialize_variational_distribution(self, prior_dist):\n        prior_prec = prior_dist.covariance_matrix.inverse()\n        prior_mean = prior_dist.mean\n        noise = torch.randn_like(prior_mean).mul_(self.mean_init_std)\n\n        self.natural_vec.data.copy_((prior_prec @ prior_mean).add_(noise))\n        self.natural_mat.data.copy_(prior_prec.mul(-0.5))\n\n\ndef _triangular_inverse(A, upper=False):\n    eye = torch.eye(A.size(-1), dtype=A.dtype, device=A.device)\n    return eye.triangular_solve(A, upper=upper).solution\n\n\ndef _phi_for_cholesky_(A):\n    \"Modifies A to be the phi function used in differentiating through Cholesky\"\n    A.tril_().diagonal(offset=0, dim1=-2, dim2=-1).mul_(0.5)\n    return A\n\n\ndef _cholesky_backward(dout_dL, L, L_inverse):\n    # c.f. https://github.com/pytorch/pytorch/blob/25ba802ce4cbdeaebcad4a03cec8502f0de9b7b3/\n    #      tools/autograd/templates/Functions.cpp\n    A = L.transpose(-1, -2) @ dout_dL\n    phi = _phi_for_cholesky_(A)\n    grad_input = (L_inverse.transpose(-1, -2) @ phi) @ L_inverse\n    # Symmetrize gradient\n    return grad_input.add(grad_input.transpose(-1, -2)).mul_(0.5)\n\n\nclass _NaturalToMuVarSqrt(torch.autograd.Function):\n    @staticmethod\n    def _forward(nat_mean, nat_covar):\n        try:\n            L_inv = torch.cholesky(-2.0 * nat_covar, upper=False)\n        except RuntimeError as e:\n            if str(e).startswith(\"cholesky\"):\n                raise RuntimeError(\n                    \"Non-negative-definite natural covariance. You probably \"\n                    \"updated it using an optimizer other than gpytorch.optim.NGD (such as Adam). \"\n                    \"This is not supported.\"\n                )\n            else:\n                raise e\n        L = _triangular_inverse(L_inv, upper=False)\n        S = L.transpose(-1, -2) @ L\n        mu = (S @ nat_mean.unsqueeze(-1)).squeeze(-1)\n        # Two choleskys are annoying, but we don't have good support for a\n        # LazyTensor of form L.T @ L\n        return mu, torch.cholesky(S, upper=False)\n\n    @staticmethod\n    def forward(ctx, nat_mean, nat_covar):\n        mu, L = _NaturalToMuVarSqrt._forward(nat_mean, nat_covar)\n        ctx.save_for_backward(mu, L)\n        return mu, L\n\n    @staticmethod\n    def _backward(dout_dmu, dout_dL, mu, L, C):\n        \"\"\"Calculate dout/d(eta1, eta2), which are:\n        eta1 = mu\n        eta2 = mu*mu^T + LL^T = mu*mu^T + Sigma\n\n        Thus:\n        dout/deta1 = dout/dmu + dout/dL dL/deta1\n        dout/deta2 = dout/dL dL/deta1\n\n        For L = chol(eta2 - eta1*eta1^T).\n        dout/dSigma = _cholesky_backward(dout/dL, L)\n        dout/deta2 = dout/dSigma\n        dSigma/deta1 = -2* (dout/dSigma) mu\n        \"\"\"\n        dout_dSigma = _cholesky_backward(dout_dL, L, C)\n        dout_deta1 = dout_dmu - 2 * (dout_dSigma @ mu.unsqueeze(-1)).squeeze(-1)\n        return dout_deta1, dout_dSigma\n\n    @staticmethod\n    def backward(ctx, dout_dmu, dout_dL):\n        \"Calculates the natural gradient with respect to nat_mean, nat_covar\"\n        mu, L = ctx.saved_tensors\n        C = _triangular_inverse(L, upper=False)\n        return _NaturalToMuVarSqrt._backward(dout_dmu, dout_dL, mu, L, C)\n", "path": "gpytorch/variational/natural_variational_distribution.py"}], "after_files": [{"content": "#!/usr/bin/env python3\n\nimport abc\n\nimport torch\n\nfrom ..distributions import MultivariateNormal\nfrom ..lazy import CholLazyTensor, TriangularLazyTensor\nfrom ..utils.cholesky import psd_safe_cholesky\nfrom ._variational_distribution import _VariationalDistribution\n\n\nclass _NaturalVariationalDistribution(_VariationalDistribution, abc.ABC):\n    r\"\"\"Any :obj:`~gpytorch.variational._VariationalDistribution` which calculates\n    natural gradients with respect to its parameters.\n    \"\"\"\n    pass\n\n\nclass NaturalVariationalDistribution(_NaturalVariationalDistribution):\n    r\"\"\"A multivariate normal :obj:`~gpytorch.variational._VariationalDistribution`,\n    parameterized by **natural** parameters.\n\n    .. note::\n       The :obj:`~gpytorch.variational.NaturalVariationalDistribution` can only\n       be used with :obj:`gpytorch.optim.NGD`, or other optimizers\n       that follow exactly the gradient direction. Failure to do so will cause\n       the natural matrix :math:`\\mathbf \\Theta_\\text{mat}` to stop being\n       positive definite, and a :obj:`~RuntimeError` will be raised.\n\n    .. seealso::\n        The `natural gradient descent tutorial\n        <examples/04_Variational_and_Approximate_GPs/Natural_Gradient_Descent.ipynb>`_\n        for use instructions.\n\n        The :obj:`~gpytorch.variational.TrilNaturalVariationalDistribution` for\n        a more numerically stable parameterization, at the cost of needing more\n        iterations to make variational regression converge.\n\n    :param int num_inducing_points: Size of the variational distribution. This implies that the variational mean\n        should be this size, and the variational covariance matrix should have this many rows and columns.\n    :param batch_shape: Specifies an optional batch size\n        for the variational parameters. This is useful for example when doing additive variational inference.\n    :type batch_shape: :obj:`torch.Size`, optional\n    :param float mean_init_std: (Default: 1e-3) Standard deviation of gaussian noise to add to the mean initialization.\n\n    \"\"\"\n\n    def __init__(self, num_inducing_points, batch_shape=torch.Size([]), mean_init_std=1e-3, **kwargs):\n        super().__init__(num_inducing_points=num_inducing_points, batch_shape=batch_shape, mean_init_std=mean_init_std)\n        scaled_mean_init = torch.zeros(num_inducing_points)\n        neg_prec_init = torch.eye(num_inducing_points, num_inducing_points).mul(-0.5)\n        scaled_mean_init = scaled_mean_init.repeat(*batch_shape, 1)\n        neg_prec_init = neg_prec_init.repeat(*batch_shape, 1, 1)\n\n        # eta1 and eta2 parameterization of the variational distribution\n        self.register_parameter(name=\"natural_vec\", parameter=torch.nn.Parameter(scaled_mean_init))\n        self.register_parameter(name=\"natural_mat\", parameter=torch.nn.Parameter(neg_prec_init))\n\n    def forward(self):\n        mean, chol_covar = _NaturalToMuVarSqrt.apply(self.natural_vec, self.natural_mat)\n        res = MultivariateNormal(mean, CholLazyTensor(TriangularLazyTensor(chol_covar)))\n        return res\n\n    def initialize_variational_distribution(self, prior_dist):\n        prior_prec = prior_dist.covariance_matrix.inverse()\n        prior_mean = prior_dist.mean\n        noise = torch.randn_like(prior_mean).mul_(self.mean_init_std)\n\n        self.natural_vec.data.copy_((prior_prec @ prior_mean.unsqueeze(-1)).squeeze(-1).add_(noise))\n        self.natural_mat.data.copy_(prior_prec.mul(-0.5))\n\n\ndef _triangular_inverse(A, upper=False):\n    eye = torch.eye(A.size(-1), dtype=A.dtype, device=A.device)\n    return eye.triangular_solve(A, upper=upper).solution\n\n\ndef _phi_for_cholesky_(A):\n    \"Modifies A to be the phi function used in differentiating through Cholesky\"\n    A.tril_().diagonal(offset=0, dim1=-2, dim2=-1).mul_(0.5)\n    return A\n\n\ndef _cholesky_backward(dout_dL, L, L_inverse):\n    # c.f. https://github.com/pytorch/pytorch/blob/25ba802ce4cbdeaebcad4a03cec8502f0de9b7b3/\n    #      tools/autograd/templates/Functions.cpp\n    A = L.transpose(-1, -2) @ dout_dL\n    phi = _phi_for_cholesky_(A)\n    grad_input = (L_inverse.transpose(-1, -2) @ phi) @ L_inverse\n    # Symmetrize gradient\n    return grad_input.add(grad_input.transpose(-1, -2)).mul_(0.5)\n\n\nclass _NaturalToMuVarSqrt(torch.autograd.Function):\n    @staticmethod\n    def _forward(nat_mean, nat_covar):\n        try:\n            L_inv = psd_safe_cholesky(-2.0 * nat_covar, upper=False)\n        except RuntimeError as e:\n            if str(e).startswith(\"cholesky\"):\n                raise RuntimeError(\n                    \"Non-negative-definite natural covariance. You probably \"\n                    \"updated it using an optimizer other than gpytorch.optim.NGD (such as Adam). \"\n                    \"This is not supported.\"\n                )\n            else:\n                raise e\n        L = _triangular_inverse(L_inv, upper=False)\n        S = L.transpose(-1, -2) @ L\n        mu = (S @ nat_mean.unsqueeze(-1)).squeeze(-1)\n        # Two choleskys are annoying, but we don't have good support for a\n        # LazyTensor of form L.T @ L\n        return mu, psd_safe_cholesky(S, upper=False)\n\n    @staticmethod\n    def forward(ctx, nat_mean, nat_covar):\n        mu, L = _NaturalToMuVarSqrt._forward(nat_mean, nat_covar)\n        ctx.save_for_backward(mu, L)\n        return mu, L\n\n    @staticmethod\n    def _backward(dout_dmu, dout_dL, mu, L, C):\n        \"\"\"Calculate dout/d(eta1, eta2), which are:\n        eta1 = mu\n        eta2 = mu*mu^T + LL^T = mu*mu^T + Sigma\n\n        Thus:\n        dout/deta1 = dout/dmu + dout/dL dL/deta1\n        dout/deta2 = dout/dL dL/deta1\n\n        For L = chol(eta2 - eta1*eta1^T).\n        dout/dSigma = _cholesky_backward(dout/dL, L)\n        dout/deta2 = dout/dSigma\n        dSigma/deta1 = -2* (dout/dSigma) mu\n        \"\"\"\n        dout_dSigma = _cholesky_backward(dout_dL, L, C)\n        dout_deta1 = dout_dmu - 2 * (dout_dSigma @ mu.unsqueeze(-1)).squeeze(-1)\n        return dout_deta1, dout_dSigma\n\n    @staticmethod\n    def backward(ctx, dout_dmu, dout_dL):\n        \"Calculates the natural gradient with respect to nat_mean, nat_covar\"\n        mu, L = ctx.saved_tensors\n        C = _triangular_inverse(L, upper=False)\n        return _NaturalToMuVarSqrt._backward(dout_dmu, dout_dL, mu, L, C)\n", "path": "gpytorch/variational/natural_variational_distribution.py"}]}