code: def forward(
      self,
      input_ids: Optional[torch.Tensor] = None,
      attention_mask: Optional[torch.Tensor] = None,
      token_type_ids: Optional[torch.Tensor] = None,
      position_ids: Optional[torch.Tensor] = None,
      head_mask: Optional[torch.Tensor] = None,
      inputs_embeds: Optional[torch.Tensor] = None,
      encoder_hidden_states: Optional[torch.Tensor] = None,
      encoder_attention_mask: Optional[torch.Tensor] = None,
      past_key_values: Optional[List[torch.FloatTensor]] = None,
      use_cache: Optional[bool] = None,
      output_attentions: Optional[bool] = None,
      output_hidden_states: Optional[bool] = None,
      return_dict: Optional[bool] = None,
  ):
      if input_ids is not None and inputs_embeds is not None:
          raise ValueError(
              "You cannot specify both input_ids and inputs_embeds at the same time"  # noqa: E501
          )
      elif input_ids is not None:
          input_shape = input_ids.size()
      elif inputs_embeds is not None:
          input_shape = inputs_embeds.size()[:-1]
      else:
          raise ValueError(
              "You have to specify either input_ids or inputs_embeds"
          )
  
      device = (
          input_ids.device
          if input_ids is not None
          else inputs_embeds.device
      )
  
      if attention_mask is None:
          attention_mask = torch.ones(input_shape, device=device)
      if token_type_ids is None:
          token_type_ids = torch.zeros(
              input_shape, dtype=torch.long, device=device
          )
  
      if self.use_ext_encoder:
          # if attention_mask.dim() == 3:
          #     extended_attention_mask = attention_mask
          # elif attention_mask.dim() == 2:
          #     extended_attention_mask = attention_mask[:, None, :].repeat(1, input_shape[1], 1) # noqa: E501
          # else:
          #     raise ValueError(
          #         "Wrong shape for input_ids (shape {}) or attention_mask (shape {})".format(# noqa: E501
          #             input_shape, attention_mask.shape
          #         )
          #     )
          assert attention_mask.dim() == 2
          extended_attention_mask = attention_mask.view(
              -1, 1, 1, attention_mask.size(-1)
          )
          m_2 = extended_attention_mask.transpose(-1, -2)
          extended_attention_mask = extended_attention_mask * m_2
          extended_attention_mask = extended_attention_mask.to(
              dtype=next(self.parameters()).dtype
          )  # fp16 compatibility
          seq_lens = torch.sum(
              attention_mask, 1, dtype=torch.int32
          ).cuda()
      else:
          # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] # noqa: E501
          # ourselves in which case we just need to make it broadcastable to all heads. # noqa: E501
          if attention_mask.dim() == 3:
              extended_attention_mask = attention_mask[:, None, :, :]
          elif attention_mask.dim() == 2:
              extended_attention_mask = attention_mask[:, None, None, :]
          else:
              raise ValueError(
                  "Wrong shape for input_ids (shape {}) or attention_mask (shape {})".format(  # noqa: E501
                      input_shape, attention_mask.shape
                  )
              )
          # Since attention_mask is 1.0 for positions we want to attend
          # and 0.0 for masked positions, this operation will create a
          # tensor which is 0.0 for positions we want to attend
          # and -10000.0 for masked positions.
          # Since we are adding it to the raw scores before the softmax,
          # this is effectively the same as removing these entirely.
          extended_attention_mask = extended_attention_mask.to(
              dtype=next(self.parameters()).dtype
          )  # fp16 compatibility
          extended_attention_mask = (
              1.0 - extended_attention_mask
          ) * -10000.0
  
      embedding_output = self.embeddings(
          input_ids=input_ids,
          position_ids=position_ids,
          token_type_ids=token_type_ids,
          inputs_embeds=inputs_embeds,
      )
      if self.use_ext_encoder:
          encoder_outputs = self.encoder(
              embedding_output, extended_attention_mask, seq_lens
          )
      else:
          head_mask = [None] * self.config.num_hidden_layers
          encoder_outputs = self.encoder(
              embedding_output,
              attention_mask=extended_attention_mask,
              head_mask=head_mask,
          )
  
      sequence_output = encoder_outputs[0]
      pooled_output = self.pooler(sequence_output)
  
      outputs = (
          sequence_output,
          pooled_output,
      ) + encoder_outputs[
          1:
      ]  # add hidden_states and attentions if they are here
      return outputs
  
imports: from typing import List, Optional
  from loguru import logger
  from nebullvm.optional_modules.torch import torch, torch_distributed
  from nebullvm.optional_modules.huggingface import BertConfig, BertEmbeddings, BertEncoder, BertPooler, BertPreTrainedModel
  from .checkpoint_quantization import checkpoint_quantization
label_desc: Use image node_id 2 for calling the BertModel obj's underlying member method code with example usage: obj.forward(input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, past_key_values, use_cache, output_attentions, output_hidden_states, return_dict) and returns: outputs