torch numberrdss

app.py

@@ -117,7 +117,8 @@ class TextGeneration:
             torch_dtype=DTYPE, low_cpu_mem_usage=False if DEVICE == "cpu" else True
         ).to(device=DEVICE, non_blocking=True)
         _ = self.model.eval()
-        device_number = -1 if DEVICE == "cpu" else int(DEVICE.split(":")[-1])
+        # -1 if DEVICE == "cpu" else int(DEVICE.split(":")[-1])
+        device_number = torch.cuda.current_device()
         self.generator = pipeline(
             self.task, model=self.model, tokenizer=self.tokenizer, device=device_number)
         print("Done")