[Bug]: Error with loading state dict for SingleTaskVariationalGP

[roussel-ryan]

What happened?

Loading a state dict into SingleTaskVariationalGP fails with the following error AttributeError: 'SingleTaskVariationalGP' object has no attribute 'train_targets'. See minimal example

Please provide a minimal, reproducible example of the unexpected behavior.

from botorch import fit_gpytorch_mll
import torch
from botorch.models import SingleTaskVariationalGP
from gpytorch.mlls import VariationalELBO

# Synthetic 1D regression data.
train_X = torch.rand(128, 1)
noise = 0.05 * torch.randn(128, 1)
train_Y = torch.sin(2 * torch.pi * train_X) + noise

model = SingleTaskVariationalGP(train_X=train_X, train_Y=train_Y)
mll = VariationalELBO(
    likelihood=model.likelihood,
    model=model.model,
    num_data=train_X.shape[-2],
)
fit_gpytorch_mll(mll)

state_dict = model.state_dict()

restored_model = SingleTaskVariationalGP(train_X=train_X, train_Y=train_Y)
restored_model.load_state_dict(state_dict)

Please paste any relevant traceback/logs produced by the example provided.

File "C:\ProgramData\mambaforge\envs\xopt-dev\Lib\site-packages\botorch\models\gpytorch.py", line 362, in load_state_dict
    untransformed_Y, untransformed_Yvar = self._extract_targets_and_noise()
                                          ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^
  File "C:\ProgramData\mambaforge\envs\xopt-dev\Lib\site-packages\botorch\models\gpytorch.py", line 304, in _extract_targets_and_noise
    Y, Yvar = extract_targets_and_noise_single_output(self)
              ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^
  File "C:\ProgramData\mambaforge\envs\xopt-dev\Lib\site-packages\botorch\models\utils\assorted.py", line 472, in extract_targets_and_noise_single_output
    Y = model.train_targets.unsqueeze(-1)
        ^^^^^^^^^^^^^^^^^^^
  File "C:\ProgramData\mambaforge\envs\xopt-dev\Lib\site-packages\torch\nn\modules\module.py", line 1965, in __getattr__
    raise AttributeError(
        f"'{type(self).__name__}' object has no attribute '{name}'"
    )
AttributeError: 'SingleTaskVariationalGP' object has no attribute 'train_targets'

BoTorch Version

0.17.2

Python Version

No response

Operating System

No response

(Optional) Describe any potential fixes you've considered to the issue outlined above.

No response

Pull Request

None

Code of Conduct

• I agree to follow BoTorch's Code of Conduct

[esantorella]

Thanks for reporting this and for the clean repro! I put up #3251 to address this.