Fix a bug in _fit_multioutput_independent that failed mll comparison (#1455)

Summary:
Pull Request resolved: #1455

The model of the repacked mll was in `eval` mode while the `repacked_mll` itself was in `train` mode, leading to the loss of `mll` and `repacked_mll` evaluating differently, thus failing the model fitting. Thankfully codecov caught it!

Reviewed By: j-wilson

Differential Revision: D40477774

fbshipit-source-id: 8f715c2bfa4b418d9c34153f184f4b779c7319b1

Author: saitcakmak

botorch/fit.py

@@ -338,7 +338,7 @@ def _fit_multioutput_independent(
             unpacked_mll = fit_gpytorch_mll(unpacked_mll, **kwargs)
 
             # Repackage submodels and copy over state_dict
-            repacked_model = model_list_to_batched(unpacked_mll.model)
+            repacked_model = model_list_to_batched(unpacked_mll.model.train())
             repacked_mll = type(mll)(repacked_model.likelihood, repacked_model)
             with state_rollback_ctx(mll, device=device("cpu")) as ckpt:
                 mll.load_state_dict(repacked_mll.state_dict())

test/test_fit.py

@@ -5,6 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import math
+import warnings
 from contextlib import nullcontext
 from copy import deepcopy
 from itertools import product
@@ -16,7 +17,7 @@
 import torch
 from botorch import fit
 from botorch.exceptions.errors import ModelFittingError, UnsupportedError
-from botorch.exceptions.warnings import OptimizationWarning
+from botorch.exceptions.warnings import BotorchWarning, OptimizationWarning
 from botorch.fit import fit_gpytorch_mll
 from botorch.models import FixedNoiseGP, HeteroskedasticSingleTaskGP, SingleTaskGP
 from botorch.models.converter import batched_to_model_list
@@ -412,7 +413,10 @@ def _test_main(self, mll, ckpt):
             return
 
         optimizer = MockOptimizer()
-        with state_rollback_ctx(mll, checkpoint=ckpt), debug(True):
+        with state_rollback_ctx(mll, checkpoint=ckpt), debug(
+            True
+        ), warnings.catch_warnings(record=True) as ws:
+            warnings.simplefilter("always", BotorchWarning)
             try:
                 fit._fit_multioutput_independent(
                     mll,
@@ -425,6 +429,7 @@ def _test_main(self, mll, ckpt):
             except Exception:
                 pass  # exception handling tested separately
             else:
+                self.assertEqual(len(ws), 0)  # Model repacking did not fail.
                 self.assertFalse(mll.training)
                 self.assertEqual(optimizer.call_count, mll.model.num_outputs)
                 self.assertTrue(
@@ -519,8 +524,13 @@ def test_fit_with_converter(self):
             with mock.patch(
                 f"{fit_gpytorch_mll.__module__}.batched_to_model_list",
                 wraps=batched_to_model_list,
-            ) as wrapped_converter:
+            ) as wrapped_converter, warnings.catch_warnings(record=True) as ws:
+                warnings.simplefilter("always", BotorchWarning)
                 fit_gpytorch_mll(mll)
+            # Check that MLL repacking succeeded.
+            self.assertFalse(
+                any("Training loss of repacked model" in str(w.message) for w in ws)
+            )
             wrapped_converter.assert_called_once()
             self.assertFalse(torch.allclose(intf.mins, torch.zeros(1, 2, **tkwargs)))
             self.assertFalse(torch.allclose(intf.ranges, torch.ones(1, 2, **tkwargs)))