Merge branch 'main_retriver_merge' into 'main'

        assert args.DDP_impl == 'local'

    # get model without FP16 and/or TorchDDP wrappers

    model = unwrap_model(model)

    for module in model:

        if args.iteration == 0 and hasattr(module,

                                           'init_state_dict_from_bert'):

            print("Initializing ICT from pretrained BERT model", flush=True)

            module.init_state_dict_from_bert()

    if args.iteration == 0 and len(unwrapped_model) == 1 \

        and hasattr(unwrapped_model[0], 'init_state_dict_from_bert'):

        print_rank_0("Initializing ICT from pretrained BERT model")

        unwrapped_model[0].init_state_dict_from_bert()

        if args.fp16:

            optimizer.reload_model_params()