Update pipeline.py

pipeline.py

@@ -3,6 +3,7 @@ from typing import Dict, List, Any
 import itertools
 from nltk import sent_tokenize
 import nltk
+import torch
 
 class PreTrainedPipeline():
 
@@ -16,8 +17,7 @@ class PreTrainedPipeline():
         self.tokenizer = AutoTokenizer.from_pretrained(path)
         
         self.model_type="t5"
-#        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.device =  "cpu"
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
     def __call__(self, inputs: str):