'GreedySearchEncoderDecoderOutput' object has no attribute 'cpu'

[doussa2504]

I am replicating the tutorial 4.1_all_tasks_are_generation.py and trying to apply it in a case of binary classification. In the model generate, the output_sentence contains only the label, but I need also the probability for each class. I tried to add "return_dict_in_generate": True and "output_scores": True to the generation_arguments, but I got the following error AttributeError: 'GreedySearchEncoderDecoderOutput' object has no attribute 'cpu' Any idea how to fix that ?