code: def init_bert_params_(module):
      """
      Initialize the weights specific to the BERT Model.
      """
  
      def normal_(data):
          # with FSDP, module params will be on CUDA, so we cast them back to CPU
          # so that the RNG is consistent with and without FSDP
          data.copy_(data.cpu().normal_(mean=0.0, std=0.02))
  
      if isinstance(module, nn.Linear):
          normal_(module.weight.data)
          if module.bias is not None:
              module.bias.data.zero_()
      if isinstance(module, nn.Embedding):
          normal_(module.weight.data)
          if module.padding_idx is not None:
              module.weight.data[module.padding_idx].zero_()
      if isinstance(module, nn.MultiheadAttention):
          normal_(module.in_proj_weight.data)
  
file_path: fairscale/tests/nn/data_parallel/test_fsdp_apply.py
imports: import functools
  import unittest
  from parameterized import parameterized
  import pytest
  import torch.nn
  from fairscale.internal import torch_version
  from .test_fsdp import CONFIG_OPTIONS, DistributedTest, NestedWrappedModule, TransformerWithSharedParams, rename_test, spawn_and_init
label_desc: Use image node_id 2 for calling a global function with example usage: init_bert_params_(module) without return types