def METHOD_NAME(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()