class_docstr: A fully Bayesian multi-task GP model with the SAAS prior.
  
  This model assumes that the inputs have been normalized to [0, 1]^d and that the
  output has been stratified standardized to have zero mean and unit variance for
  each task.The SAAS model [Eriksson2021saasbo]_ with a Matern-5/2 is used as data
  kernel by default.
  
  You are expected to use `fit_fully_bayesian_model_nuts` to fit this model as it
  isn't compatible with `fit_gpytorch_model`.
  
  Example:
      >>> X1, X2 = torch.rand(10, 2), torch.rand(20, 2)
      >>> i1, i2 = torch.zeros(10, 1), torch.ones(20, 1)
      >>> train_X = torch.cat([
      >>>     torch.cat([X1, i1], -1), torch.cat([X2, i2], -1),
      >>> ])
      >>> train_Y = torch.cat(f1(X1), f2(X2)).unsqueeze(-1)
      >>> train_Yvar = 0.01 * torch.ones_like(train_Y)
      >>> mtsaas_gp = SaasFullyBayesianMultiTaskGP(
      >>>     train_X, train_Y, train_Yvar, task_feature=-1,
      >>> )
      >>> fit_fully_bayesian_model_nuts(mtsaas_gp)
      >>> posterior = mtsaas_gp.posterior(test_X)
docstr: Custom logic for loading the state dict.
  
  The standard approach of calling `load_state_dict` currently doesn't play well
  with the `SaasFullyBayesianMultiTaskGP` since the `mean_module`, `covar_module`
  and `likelihood` aren't initialized until the model has been fitted. The reason
  for this is that we don't know the number of MCMC samples until NUTS is called.
  Given the state dict, we can initialize a new model with some dummy samples and
  then load the state dict into this model. This currently only works for a
  `MultitaskSaasPyroModel` and supporting more Pyro models likely requires moving
  the model construction logic into the Pyro model itself.
  
  TODO: If this were to inherif from `SaasFullyBayesianSingleTaskGP`, we could
  simplify this method and eliminate some others.
code: def load_state_dict(
      self, state_dict: Mapping[str, Any], strict: bool = True
  ):
      r"""Custom logic for loading the state dict.
  
      The standard approach of calling `load_state_dict` currently doesn't play well
      with the `SaasFullyBayesianMultiTaskGP` since the `mean_module`, `covar_module`
      and `likelihood` aren't initialized until the model has been fitted. The reason
      for this is that we don't know the number of MCMC samples until NUTS is called.
      Given the state dict, we can initialize a new model with some dummy samples and
      then load the state dict into this model. This currently only works for a
      `MultitaskSaasPyroModel` and supporting more Pyro models likely requires moving
      the model construction logic into the Pyro model itself.
  
      TODO: If this were to inherif from `SaasFullyBayesianSingleTaskGP`, we could
      simplify this method and eliminate some others.
      """
      if not isinstance(self.pyro_model, MultitaskSaasPyroModel):
          raise NotImplementedError(  # pragma: no cover
              "load_state_dict only works for MultitaskSaasPyroModel"
          )
      raw_mean = state_dict["mean_module.raw_constant"]
      num_mcmc_samples = len(raw_mean)
      dim = (
          self.pyro_model.train_X.shape[-1] - 1
      )  # Removing 1 for the task feature.
      tkwargs = {"device": raw_mean.device, "dtype": raw_mean.dtype}
      # Load some dummy samples
      mcmc_samples = {
          "mean": torch.ones(num_mcmc_samples, **tkwargs),
          "lengthscale": torch.ones(num_mcmc_samples, dim, **tkwargs),
          "outputscale": torch.ones(num_mcmc_samples, **tkwargs),
          "task_lengthscale": torch.ones(
              num_mcmc_samples, self._rank, **tkwargs
          ),
          "latent_features": torch.ones(
              num_mcmc_samples, self.num_tasks, self._rank, **tkwargs
          ),
      }
      if self.pyro_model.train_Yvar is None:
          mcmc_samples["noise"] = torch.ones(
              num_mcmc_samples, **tkwargs
          )
      (
          self.mean_module,
          self.covar_module,
          self.likelihood,
          self.task_covar_module,
          self.latent_features,
      ) = self.pyro_model.load_mcmc_samples(mcmc_samples=mcmc_samples)
      # Load the actual samples from the state dict
      super().load_state_dict(state_dict=state_dict, strict=strict)
  
imports: from typing import Any, Dict, List, Mapping, NoReturn, Optional, Tuple, Union
  import pyro
  import torch
  from botorch.acquisition.objective import PosteriorTransform
  from botorch.models.fully_bayesian import matern52_kernel, MIN_INFERRED_NOISE_LEVEL, PyroModel, reshape_and_detach, SaasPyroModel
  from botorch.models.multitask import MultiTaskGP
  from botorch.models.transforms.input import InputTransform
  from botorch.models.transforms.outcome import OutcomeTransform
  from botorch.posteriors.fully_bayesian import GaussianMixturePosterior, MCMC_DIM
  from botorch.utils.datasets import MultiTaskDataset, SupervisedDataset
  from gpytorch.distributions.multivariate_normal import MultivariateNormal
  from gpytorch.kernels import MaternKernel
  from gpytorch.kernels.kernel import Kernel
  from gpytorch.likelihoods.likelihood import Likelihood
  from gpytorch.means.mean import Mean
  from torch import Tensor
  from torch.nn.parameter import Parameter
label_desc: Use image node_id 12 for calling the SaasFullyBayesianMultiTaskGP obj's underlying member method code with example usage: obj.load_state_dict(state_dict, strict) without return types