examples/demo_torch_multi_gpu.py/main
def main(current_gpu_index, num_gpu):
    # setup the process groups
    setup(current_gpu_index, num_gpu)

    #################################################################
    ######## Writing a torch training loop for a Keras model ########
    #################################################################

    dataset = get_data()
    model = get_model()

    # prepare the dataloader
    dataloader = prepare(dataset, current_gpu_index, num_gpu, batch_size)

    # Instantiate the torch optimizer
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    # Instantiate the torch loss function
    loss_fn = nn.CrossEntropyLoss()

    # Put model on device
    model = model.to(current_gpu_index)
    ddp_model = DDP(
        model, device_ids=[current_gpu_index], output_device=current_gpu_index
    )

    train(ddp_model, dataloader, num_epochs, optimizer, loss_fn)

    ################################################################
    ######## Using a Keras model or layer in a torch Module ########
    ################################################################

    torch_module = MyModel().to(current_gpu_index)
    ddp_torch_module = DDP(
        torch_module,
        device_ids=[current_gpu_index],
        output_device=current_gpu_index,
    )

    # Instantiate the torch optimizer
    optimizer = optim.Adam(torch_module.parameters(), lr=learning_rate)

    # Instantiate the torch loss function
    loss_fn = nn.CrossEntropyLoss()

    train(ddp_torch_module, dataloader, num_epochs, optimizer, loss_fn)

    cleanup()