code: def get_init_weights_func():
      @torch.no_grad()
      def init_weights(m):
          if isinstance(m, torch.nn.Linear):
              m.weight.fill_(1.0)
  
      return init_weights
  
file_path: apex/apex/contrib/test/optimizers/test_distributed_fused_lamb.py
imports: import os
  import inspect
  import torch
  from torch.cuda.amp import GradScaler
  from torch.testing._internal import common_utils
  from apex.parallel.distributed import flat_dist_call
  from apex.contrib.optimizers.distributed_fused_lamb import DistributedFusedLAMB
  from apex.transformer.testing.distributed_test_base import NcclDistributedTestBase
label_desc: Use image node_id 1 for calling a global function with example usage: get_init_weights_func() and returns: init_weights