Repository: cornellius-gp/gpytorch

Issue:
[Bug] Resize error thrown when using Batch Independent Multi Output
# 🐛 Bug

<!-- A clear and concise description of what the bug is. -->
Essentially I tried to take a crack at implementing the `from_batch_mvn` style for Batch Independent Multi Output GP based on the tutorial: https://docs.gpytorch.ai/en/v1.2.1/examples/03_Multitask_Exact_GPs/Batch_Independent_Multioutput_GP.html

Training works just fine, but when I try to predict in `eval()` mode, I get a resize error inside of gpytorch code.
It is interesting to note that when my output_dim is 1, I do not face any issues.

My code below is not really reproducible. I'm using [pytorch-lightning](https://www.pytorchlightning.ai/) to take care of the training loops and optimization calls. Its pretty neat and was happy I got it working with GPyTorch, maybe I'll submit a PR with an example usage of gpytorch and pytorch lightning together. In any case I'll add a minimal working example soon.

## To reproduce

** Code snippet to reproduce **
```python
"""Gaussian Process Model."""

import argparse

import gpytorch

import pytorch_lightning as lt

import torch


class BIMOEGP(gpytorch.models.ExactGP):
    """batch independent multioutput exact gp model."""

    def __init__(self, train_input_data, train_output_data, likelihood):
        """Initialize gp model with mean and covar."""
        super().__init__(train_input_data, train_output_data, likelihood)

        output_dim = train_output_data.shape[1]
        output_dim_torch = torch.Size([output_dim])

        self.mean_module = \
            gpytorch.means.ConstantMean(batch_shape=output_dim_torch)

        self.covar_module = gpytorch.kernels.ScaleKernel(
            gpytorch.kernels.RBFKernel(batch_shape=output_dim_torch),
            batch_shape=output_dim_torch)

    # pylint: disable=arguments-differ
    def forward(self, input_):
        """Compute prediction."""
        mean = self.mean_module(input_)
        covar = self.covar_module(input_)

        return \
            gpytorch.distributions.MultitaskMultivariateNormal.from_batch_mvn(
                gpytorch.distributions.MultivariateNormal(mean, covar))


# pylint: disable=too-many-ancestors
class BIMOEGPModel(lt.core.lightning.LightningModule):
    """batch independent multioutput exact gp model."""

    def __init__(self, hparams, train_input_data, train_output_data):
        """Initialize gp model with mean and covar."""
        super().__init__()

        self.hparams = hparams

        output_dim = train_output_data.shape[1]
        self.likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(
            num_tasks=output_dim)

        self.bimoegp = BIMOEGP(train_input_data, train_output_data,
                               self.likelihood)

        self.mll = gpytorch.mlls.ExactMarginalLogLikelihood(
            self.likelihood, self.bimoegp)

    # pylint: disable=arguments-differ
    def forward(self, input_):
        """Compute prediction."""
        return self.bimoegp(input_)

    # pylint: disable=unused-argument
    def training_step(self, batch, batch_idx):
        """Compute training loss."""
        input_, target = batch
        output = self(input_)

        loss = -self.mll(output, target)

        return {'loss': loss}

    def configure_optimizers(self):
        """Create optimizer."""
        optimizer = torch.optim.Adam(
            self.parameters(),
            lr=self.hparams.learning_rate)

        return optimizer

    # pylint: disable=unused-argument
    def validation_step(self, batch, batch_idx):
        """Compute validation loss."""
        input_, target = batch
        output = self(input_)

        loss = -self.mll(output, target)

        return {'val_loss': loss}

    # pylint: disable=no-self-use
    def validation_epoch_end(self, outputs):
        """Record validation loss."""
        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
        self.log('avg_val_loss', avg_loss)

    # pylint: disable=unused-argument
    def test_step(self, batch, batch_idx):
        """Compute testing loss."""
        input_, target = batch
        output = self(input_)

        loss = -self.mll(output, target)

        return {'test_loss': loss}

    def test_epoch_end(self, outputs):
        """Record average test loss."""
        avg_loss = torch.stack([x['test_loss'] for x in outputs]).mean()
        self.log('avg_test_loss', avg_loss)

    @ staticmethod
    def add_model_specific_args(parent_parser):
        """Parse model specific hyperparameters."""
        parser = argparse.ArgumentParser(
            parents=[parent_parser], add_help=False)
        parser.add_argument('--learning_rate', type=float, default=1e-3)

        return parser
```

** Stack trace/error message **
```
  File "/usr/lib/python3.8/site-packages/pl_utils/models/gp.py", line 73, in forward
    return self.bimoegp(input_)
  File "/usr/lib/python3.8/site-packages/gpytorch/models/exact_gp.py", line 283, in __call__
    self.prediction_strategy = prediction_strategy(
  File "/usr/lib/python3.8/site-packages/gpytorch/models/exact_prediction_strategies.py", line 33, in prediction_strategy
    return cls(train_inputs, train_prior_dist, train_labels, likelihood)
  File "/usr/lib/python3.8/site-packages/gpytorch/models/exact_prediction_strategies.py", line 40, in __init__
    train_labels = train_labels.view(*train_labels.shape[: -len(train_shape)], train_shape.numel())
RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.
```

## Expected Behavior

<!-- A clear and concise description of what you expected to happen. -->
There should be no error related to resizing.

## System information

**Please complete the following information:**
- GPyTorch Version: 1.2.1
- PyTorch Version: 1.7.0
- Arch Linux

## Additional context
Downstream issue: https://github.com/acxz/pl-utils/issues/7



Discussion:

Resolving patch:
diff --git a/gpytorch/models/exact_prediction_strategies.py b/gpytorch/models/exact_prediction_strategies.py
--- a/gpytorch/models/exact_prediction_strategies.py
+++ b/gpytorch/models/exact_prediction_strategies.py
@@ -37,7 +37,7 @@ class DefaultPredictionStrategy(object):
     def __init__(self, train_inputs, train_prior_dist, train_labels, likelihood, root=None, inv_root=None):
         # Flatten the training labels
         train_shape = train_prior_dist.event_shape
-        train_labels = train_labels.view(*train_labels.shape[: -len(train_shape)], train_shape.numel())
+        train_labels = train_labels.reshape(*train_labels.shape[: -len(train_shape)], train_shape.numel())
 
         self.train_inputs = train_inputs
         self.train_prior_dist = train_prior_dist


Test patch: