Function main

generate_samples.py:295–323

Main training program.

()

def main():
    """Main training program."""

    print('Generate Samples')

    # Disable CuDNN.
    torch.backends.cudnn.enabled = False

    # Arguments.
    args = get_args()

    # Pytorch distributed.
    initialize_distributed(args)

    # set device, this args.device is only used in inference
    if args.device is not None:
        device = int(args.device)
        torch.cuda.set_device(device)

    # Random seeds for reproducability.
    set_random_seed(args.seed)

    # get the tokenizer
    tokenizer = prepare_tokenizer(args)

    # Model, optimizer, and learning rate.
    model = setup_model(args)

    generate_images_continually(model, args)

if __name__ == "__main__":
    main()