code: def model_init_and_apply_custom_weight_init(
      model_init_fn, *args, **kwargs
  ):
      model = model_init_fn(*args, **kwargs)
      model.apply(init_bert_params_)
      return model
  
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 1 for calling a global function with example usage: model_init_and_apply_custom_weight_init(model_init_fn) and returns: model