use device_number for setting GPU

avoid error: Expected a torch.device with a specified index or an integer, but got:cuda

handler.py

@@ -18,15 +18,18 @@ class EndpointHandler():
     SENTENCE_TRANSFORMER_MODEL_NAME = "multi-qa-mpnet-base-dot-v1"
     QUESTION_ANSWER_MODEL_NAME = "vblagoje/bart_lfqa"
     SUMMARIZER_MODEL_NAME = "philschmid/bart-large-cnn-samsum"
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    device_number = 0 if torch.cuda.is_available() else -1
 
     def __init__(self, path=""):
 
         device = "cuda" if torch.cuda.is_available() else "cpu"
         device_number = 0 if torch.cuda.is_available() else -1
         print(f'whisper and question_answer_model will use: {device}')
+        print(f'whisper and question_answer_model will use device_number: {device_number}')
 
         t0 = time.time()
-        self.whisper_model = whisper.load_model(self.WHISPER_MODEL_NAME).to(device)
+        self.whisper_model = whisper.load_model(self.WHISPER_MODEL_NAME).to(device_number)
         t1 = time.time()
 
         total = t1 - t0
@@ -45,10 +48,11 @@ class EndpointHandler():
 
         total = t1 - t0
         print(f'Finished loading summarizer in {total} seconds')
-        
+
         self.question_answer_tokenizer = AutoTokenizer.from_pretrained(self.QUESTION_ANSWER_MODEL_NAME)
         t0 = time.time()
-        self.question_answer_model = AutoModelForSeq2SeqLM.from_pretrained(self.QUESTION_ANSWER_MODEL_NAME).to(device)
+        self.question_answer_model = AutoModelForSeq2SeqLM.from_pretrained \
+            (self.QUESTION_ANSWER_MODEL_NAME).to(device_number)
         t1 = time.time()
         total = t1 - t0
         print(f'Finished loading question_answer_model in {total} seconds')
@@ -199,22 +203,25 @@ class EndpointHandler():
         conditioned_doc = "<P> " + " <P> ".join([d for d in documents])
         query_and_docs = "question: {} context: {}".format(query, conditioned_doc)
 
-        model_input = self.question_answer_tokenizer(query_and_docs, truncation=False, padding=True, return_tensors="pt")
-
-        generated_answers_encoded = self.question_answer_model.generate(input_ids=model_input["input_ids"].to(self.device),
-                                                attention_mask=model_input["attention_mask"].to(self.device),
-                                                min_length=64,
-                                                max_length=256,
-                                                do_sample=False, 
-                                                early_stopping=True,
-                                                num_beams=8,
-                                                temperature=1.0,
-                                                top_k=None,
-                                                top_p=None,
-                                                eos_token_id=self.question_answer_tokenizer.eos_token_id,
-                                                no_repeat_ngram_size=3,
-                                                num_return_sequences=1)
-        answer = self.question_answer_tokenizer.batch_decode(generated_answers_encoded, skip_special_tokens=True,clean_up_tokenization_spaces=True)
+        model_input = self.question_answer_tokenizer(query_and_docs, truncation=False, padding=True,
+                                                     return_tensors="pt")
+
+        generated_answers_encoded = self.question_answer_model.generate(
+            input_ids=model_input["input_ids"].to(self.device),
+            attention_mask=model_input["attention_mask"].to(self.device),
+            min_length=64,
+            max_length=256,
+            do_sample=False,
+            early_stopping=True,
+            num_beams=8,
+            temperature=1.0,
+            top_k=None,
+            top_p=None,
+            eos_token_id=self.question_answer_tokenizer.eos_token_id,
+            no_repeat_ngram_size=3,
+            num_return_sequences=1)
+        answer = self.question_answer_tokenizer.batch_decode(generated_answers_encoded, skip_special_tokens=True,
+                                                             clean_up_tokenization_spaces=True)
         return answer
 
     @staticmethod