Function get_pp_module

training/pipeline_parallel/dist_pp_utils.py:4–10

(args, config, device, use_dp)

def get_pp_module(args, config, device, use_dp):
    
    if args.pp_mode == 'gpipe':
        return GpipeAsync(args, config, device, use_dp)
    else:
        print("Not recognize this pipeline parallel mode.")
        assert False