Fixed TextGenerationPipeline on torch + GPU (#5629)

* Pytorch gpu => cpu proper device * Memoryless XLNet warning + fixed memories during generation * Revert "Memoryless XLNet warning + fixed memories during generation" This reverts commit 3d3251ff * Took the operations on the generated_sequence out of the ensure_device scope
2020-07-09 16:29:32 -04:00
parent 760f726e51
commit 02a0b43014
1 changed files with 2 additions and 0 deletions
@@ -689,6 +689,8 @@ class TextGenerationPipeline(Pipeline):

            result = []
            for generated_sequence in output_sequences:
+                if self.framework == "pt" and generated_sequence is not None:
+                    generated_sequence = generated_sequence.cpu()
                generated_sequence = generated_sequence.numpy().tolist()
                record = {}
                if return_tensors: