Merge pull request #2382 from fjzzq2002/master

Fix training status of noise model of `HeteroskedasticNoise` after exceptions

Author: Balandat

gpytorch/likelihoods/noise_models.py

@@ -125,13 +125,15 @@ def forward(
         if noise is not None:
             return DiagLinearOperator(noise)
         training = self.noise_model.training  # keep track of mode
-        self.noise_model.eval()  # we want the posterior prediction of the noise model
-        with settings.detach_test_caches(False), settings.debug(False):
-            if len(params) == 1 and not torch.is_tensor(params[0]):
-                output = self.noise_model(*params[0])
-            else:
-                output = self.noise_model(*params)
-        self.noise_model.train(training)
+        try:
+            self.noise_model.eval()  # we want the posterior prediction of the noise model
+            with settings.detach_test_caches(False), settings.debug(False):
+                if len(params) == 1 and not torch.is_tensor(params[0]):
+                    output = self.noise_model(*params[0])
+                else:
+                    output = self.noise_model(*params)
+        finally:
+            self.noise_model.train(training)
         if not isinstance(output, MultivariateNormal):
             raise NotImplementedError("Currently only noise models that return a MultivariateNormal are supported")
         # note: this also works with MultitaskMultivariateNormal, where this

test/likelihoods/test_noise_models.py

@@ -0,0 +1,38 @@
+#!/usr/bin/env python3
+
+import unittest
+
+import torch
+
+import gpytorch
+from gpytorch.likelihoods import HeteroskedasticNoise
+
+
+class NumericallyUnstableModelExample(gpytorch.models.GP):
+    def __init__(self):
+        super(NumericallyUnstableModelExample, self).__init__()
+        self.fail_arithmetic = False
+
+    def train(self, mode=True):
+        if mode:
+            self.fail_arithmetic = False  # reset on .train()
+        super().train(mode=mode)
+
+    def forward(self, x):
+        if self.fail_arithmetic:
+            raise ArithmeticError()
+        return gpytorch.distributions.MultivariateNormal(torch.tensor([-3.0]), torch.tensor([[2.0]]))
+
+
+class TestNoiseModels(unittest.TestCase):
+    def test_heteroskedasticnoise_error(self):
+        noise_model = NumericallyUnstableModelExample().to(torch.double)
+        likelihood = HeteroskedasticNoise(noise_model)
+        self.assertEqual(noise_model.training, True)
+        self.assertEqual(likelihood.training, True)
+        noise_model.fail_arithmetic = True
+        test_x = torch.tensor([[3.0]])
+        with self.assertRaises(ArithmeticError):
+            likelihood(test_x)
+        self.assertEqual(likelihood.training, True)
+        likelihood(test_x)