class_docstr: Allocates the covariance matrix on distributed devices, e.g. multiple GPUs.
  
  Args:
      base_kernel: Base kernel to distribute
      device_ids: list of `torch.device` objects to place kernel chunks on
      output_device: Device where outputs will be placed
class_docstr_tok: ["Allocates","the","covariance","matrix","on","distributed","devices",",","e.g",".","multiple","GPUs",".","Args",":","base_kernel",":","Base","kernel","to","distribute","device_ids",":","list","of","`","torch.device","`","objects","to","place","kernel","chunks","on","output_device",":","Device","where","outputs","will","be","placed"]
code: def forward(self, x1, x2, diag=False, **kwargs):
      if diag:
          return self.module.forward(x1, x2, diag=True, **kwargs).to(
              self.output_device
          )
  
      if x1.size(-2) < len(self.device_ids) + 1:
          return self.module.forward(x1, x2, diag=diag, **kwargs).to(
              self.output_device
          )
  
      if not x1.device == self.__cached_x1.device or not torch.equal(
          x1, self.__cached_x1
      ):
          self._x1_scattered, self._kwargs = self.scatter(
              (x1,), kwargs, self.device_ids
          )
          self.__cached_x1 = x1
  
      if not x2.device == self.__cached_x2.device or not torch.equal(
          x2, self.__cached_x2
      ):
          self._x2_subs = [
              x2.to(x1_[0].device) for x1_ in self._x1_scattered
          ]
          self.__cached_x2 = x2
  
      inputs = tuple(
          (x1_[0], x2_)
          for x1_, x2_ in zip(self._x1_scattered, self._x2_subs)
      )
  
      if not self.device_ids:
          return self.module.forward(*inputs, **self._kwargs)
  
      if len(self.device_ids) == 1:
          return self.module.forward(*inputs[0], **self._kwargs[0])
  
      # JIT modules can't be pickled and replicated yet
      # But reinitializing the distance_module every forward pass
      # is slow and should be removed once JIT modules can be pickled
      def set_distance_module_to_none(module):
          if hasattr(module, "distance_module"):
              module.distance_module = None
  
      self.module.apply(set_distance_module_to_none)
      # Can't cache the replication because the base kernel module can change every time (e.g. param updates)
      replicas = self.replicate(
          self.module, self.device_ids[: len(inputs)]
      )
  
      # TODO: parallel_apply might be too heavyweight in some cases?
      with settings.lazily_evaluate_kernels(False):
          outputs = self.parallel_apply(replicas, inputs, self._kwargs)
  
      return self.gather(outputs, self.output_device)
  
file_path: gpytorch/gpytorch/kernels/multi_device_kernel.py
imports: from typing import List, Optional
  import torch
  from linear_operator import to_linear_operator
  from linear_operator.operators import CatLinearOperator
  from torch.nn.parallel import DataParallel
  from ..None import settings
  from .kernel import Kernel
label_desc: Use image node_id 3 for calling the MultiDeviceKernel obj's underlying member method code with example usage: obj.forward(x1, x2, diag) and returns: self, self, self, self, self