code: def build_model(
      model_fn,
      enable_fused_normalization=False,
      enable_tensor_parallelism=False,
      enable_flash_attention=False,
      enable_jit_fused=False,
  ):
      # create new model
      org_model = model_fn()
  
      # shard model
      shard_config = ShardConfig(
          enable_fused_normalization=enable_fused_normalization,
          enable_tensor_parallelism=enable_tensor_parallelism,
          enable_flash_attention=enable_flash_attention,
          enable_jit_fused=enable_jit_fused,
      )
      model_copy = copy.deepcopy(org_model)
      shard_former = ShardFormer(shard_config=shard_config)
      sharded_model, shared_params = shard_former.optimize(model_copy)
      return org_model.cuda(), sharded_model.cuda()
  
label_desc: Use image node_id 1 for calling a global function with example usage: build_model(model_fn, enable_fused_normalization, enable_tensor_parallelism, enable_flash_attention, enable_jit_fused) and returns: org_model,  sharded_model