code: def module_policy(self):
      policy = super().module_policy()
      policy = self.add_lm_head_policy(policy)
      policy = self.add_lm_prediction_policy(policy)
      from transformers.models.bert.modeling_bert import (
          BertForPreTraining,
      )
  
      if self.pipeline_stage_manager:
          self.set_pipeline_forward(
              model_cls=BertForPreTraining,
              new_forward=BertPipelineForwards.bert_for_pretraining_forward,
              policy=policy,
          )
      return policy
  
imports: from functools import partial
  from typing import Callable, Dict, List
  import torch.nn
  from torch import Tensor
  from torch.nn import Module
  import colossalai.shardformer.layer
  from ..modeling.bert import BertPipelineForwards, bert_sequence_parallel_forward_fn, get_bert_flash_attention_forward, get_jit_fused_bert_output_forward, get_jit_fused_bert_self_output_forward
  from ..modeling.jit import get_jit_fused_dropout_add_func
  from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
label_desc: Use image node_id 1 for calling the BertForPreTrainingPolicy obj's underlying member method code with example usage: obj.module_policy() and returns: policy