I was wondering how one might go about generating text on a machine with multiple GPU? Here is some example code:
import pandas as pd
from transformers import GPT2Tokenizer, GPT2LMHeadModel

tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2LMHeadModel.from_pretrained('gpt2')

def generate_text(primer_sent: str, chunk_sent_length: int, is_finetuning_current_model=None):
    encoded_results = tokenizer(primer_sent, return_tensors='pt')
    tmp_ids = encoded_results['input_ids']
    max_length = tmp_ids.shape[1] + chunk_sent_length
    gen_ids = model.generate(input_ids=tmp_ids,
                                  max_length=max_length,
                                  is_finetuning_current_model=is_finetuning_current_model,
                                  do_sample=True,
                                  top_k=50,
                                  top_p=0.95,
                                  pad_token_id=50256)
    decoded_str = tokenizer.decode(gen_ids, skip_special_tokens=True)
    return decoded_str



def generate_text_samples(data_filename: str, chunk_sent_length: int):
    sent_df = pd.read_csv(data_filename)
    sent_df['gen'] = sent_df[0].apply(generate_text, args=(chunk_sent_length,))
    return sent_df

text_df = generate_text_samples('/path/to/data.csv', 50)

When I run this code on my machine, it appears that it is all being done on the CPU. Is it possible to distribute this workflow across the multiple gpu I have? Thanks in advance for your help!