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 __init__(
      self,
      base_kernel: Kernel,
      device_ids: List[torch.device],
      output_device: Optional[torch.device] = None,
      create_cuda_context: Optional[bool] = True,
      **kwargs,
  ):
      # Need to warm up each GPU otherwise scattering in forward will be
      # EXTREMELY slow. This memory will be available as soon as we leave __init__
      if create_cuda_context:
          for d in device_ids:
              _ = torch.tensor([], device=d)
  
      DataParallel.__init__(
          self,
          module=base_kernel,
          device_ids=device_ids,
          output_device=output_device,
          dim=-2,
      )
  
      self.output_device = (
          output_device if output_device else device_ids[0]
      )
  
      self.__cached_x1 = torch.empty(1)
      self.__cached_x2 = torch.empty(1)
  
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 1 to create a new MultiDeviceKernel object from inherited base classes: DataParallel, Kernel with example: obj = MultiDeviceKernel(base_kernel, device_ids, output_device, create_cuda_context)