Pytorch code snippet to parallelly train on multiple GPUs

spell_classifier.py

def initialize_model(hidden_dim, hidden_layers, lr, device):
    # As in line 11, pass your model to nn.DataParallel class
    # Pytorch would now parallelize your model trainig if multiple GPU's are found during initialization.
   
    input_dim = alph_len * 3
    hidden_dim = hidden_dim
    layer_dim = hidden_layers
    output_dim = 2

    model = LSTMModel(input_dim, hidden_dim, layer_dim, output_dim, device)
    model = nn.DataParallel(model)
    model = model.to(device)
    learning_rate = lr
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    criterion = nn.CrossEntropyLoss()

    return model, criterion, optimizer