docstr: MultiheadAttention takes additional layer inputs key and value.
  
  Args:
      device: torch.device
      stats: dictionary from item names (e.g. input_tensor, module, labels)
      to their respective sizes
  
  Returns: updated stats dictionary with input tensor, key, and value size
code: def _inputs_to(
      self, device: torch.device, stats: Dict[str, int]
  ) -> Dict[str, int]:
      """MultiheadAttention takes additional layer inputs key and value.
  
      Args:
          device: torch.device
          stats: dictionary from item names (e.g. input_tensor, module, labels)
          to their respective sizes
  
      Returns: updated stats dictionary with input tensor, key, and value size
      """
      stats = super()._inputs_to(device=device, stats=stats)
      self._key = self._key.to(device)
      stats["key"] = self._get_memory_difference(
          device=device, stats=stats
      )
      self._value = self._value.to(device)
      stats["value"] = self._get_memory_difference(
          device=device, stats=stats
      )
      return stats
  
imports: from enum import Enum
  from typing import Callable, Dict, Optional, Tuple, Type, Union
  import torch
  import torch.nn
  import torch.nn.functional
  from opacus.grad_sample import GradSampleModule
  from opacus.grad_sample.utils import wrap_model
  from opacus.layers import DPGRU, DPLSTM, DPRNN, DPMultiheadAttention
  from opacus.layers.dp_rnn import DPRNNBase
label_desc: Use image node_id 2 for calling the MHABase obj's underlying member method code with example usage: obj._inputs_to(device, stats) and returns: stats