CUDA error when substituting with Pegasus-xsum

[jgriffi]

Information

The question or comment is about chapter:

• [ ] Introduction
• [ ] Text Classification
• [ ] Transformer Anatomy
• [ ] Multilingual Named Entity Recognition
• [ ] Text Generation
• [x] Summarization
• [ ] Question Answering
• [ ] Making Transformers Efficient in Production
• [ ] Dealing with Few to No Labels
• [ ] Training Transformers from Scratch
• [ ] Future Directions

CUDA error when substituting with Pegasus-xsum

RuntimeError: CUDA error: CUBLAS_STATUS_NOT_INITIALIZED when calling cublasCreate(handle)

To recreate:

device = "cuda" if torch.cuda.is_available() else "cpu"

def chunks(list_of_elements, batch_size):
    """
    Yield successive batch-sized chunks form list_of_elements.
    """
    for i in range(0, len(list_of_elements), batch_size):
        yield list_of_elements[i : i + batch_size]

def evaluate_summaries_pegasus(dataset, metric, model, tokenizer,
                               batch_size=16, device=device,
                               column_text="text",
                               column_summary="title"):
    abstract_batches = list(chunks(dataset[column_text], batch_size))
    target_batches = list(chunks(dataset[column_summary], batch_size))

    for abstract_batch, target_batch in tqdm(
        zip(abstract_batches, target_batches), total=len(abstract_batches)):

        inputs = tokenizer(abstract_batch, max_length=1024, truncation=True,
                           padding="max_length", return_tensors="pt")

        summaries = model.generate(input_ids=inputs["input_ids"].to(device),
                                   attention_mask=inputs["attention_mask"].to(device),
                                   length_penalty=0.8, num_beams=8, max_length=128)

        decoded_summaries = [tokenizer.decode(s, skip_special_tokens=True,
                                              clean_up_tokenization_spaces=True) 
                            for s in summaries]
        decoded_summaries = [d.replace("<n>", " ") for d in decoded_summaries]
        metric.add_batch(predictions=decoded_summaries,
                         references=target_batch)

    score = metric.compute()
    return score

tokenizer = AutoTokenizer.from_pretrained("google/pegasus-xsum")
pegasus_xsum_model = AutoModelForSeq2SeqLM.from_pretrained("google/pegasus-xsum").to(device)
score_pegasus_xsum = evaluate_summaries_pegasus(test_sampled, rouge_metric,
                                            pegasus_xsum_model, tokenizer, batch_size=4)
rouge_dict_pegasus_xsum = dict((rn, score_pegasus_xsum[rn].mid.fmeasure) for rn in rouge_names)
pegasus_xsum_sample = pd.DataFrame(rouge_dict_pegasus_xsum, index=["pegasus_xsum_sample"])