class_docstr: The DeepSet surrogate is similar to the MLP surrogate but makes predictions
  for ensembles rather than single models.
  
  Currently, it does not support the use of dataset features.
code: def _fit(
      self, X: List[Config[EnsembleConfig]], y: npt.NDArray[np.float32]
  ) -> None:
      # Fit transformers to infer dimensionality
      X_numpy_list = self.config_transformer.fit_transform(X)
      X_numpy = np.concatenate(X_numpy_list)
      X_lengths_numpy = np.array([x.shape[0] for x in X_numpy_list])
  
      input_dim = len(self.config_transformer.feature_names_)
      output_dim = y.shape[1]
  
      # For initializing data, we prepare group IDs for the datasets
      mapping = {d: i for i, d in enumerate({x.dataset for x in X})}
  
      # For each output variable, we need to train a separate model
      self.models_ = []
      for i in range(output_dim):
          model = self._init_model(input_dim)
          module = DeepSetLightningModule(
              model, self.loss, self.weight_decay
          )
  
          # Train on output variable i
          dataset = TensorDataset(
              torch.from_numpy(X_numpy).float(),
              torch.from_numpy(X_lengths_numpy).long(),
              torch.from_numpy(y[:, i : i + 1]).float(),
              torch.as_tensor(
                  [mapping[x.dataset] for x in X], dtype=torch.long
              ),
          )
          train_loader = DataLoader(dataset, batch_size=len(dataset))
          self._trainer.fit(module, train_dataloaders=train_loader)
  
          # Add to models
          self.models_.append(model)
  
imports: from typing import cast, List, Literal, Optional
  import numpy
  import numpy.typing
  import pytorch_lightning
  import torch
  from lightkit.data import DataLoader
  from torch import nn
  from torch.utils.data import TensorDataset
  from tsbench.config import Config, EnsembleConfig
  from tsbench.evaluations.tracking import EnsembleTracker
  from tsbench.surrogate.torch.deepset import DeepSetModel
  from ._base import OutputNormalization, Surrogate
  from ._factory import register_ensemble_surrogate
  from .torch import DeepSetLightningModule, ListMLELoss
  from .transformers import EnsembleConfigTransformer
label_desc: Use image node_id 3 for calling the DeepSetSurrogate obj's underlying member method code with example usage: obj._fit(X, y) without return types