From cc1d938e0fd0b36a6aba15e3903f63882beee10a Mon Sep 17 00:00:00 2001 From: drbh Date: Fri, 8 Mar 2024 15:33:37 -0500 Subject: [PATCH] fix: prefer `device_id` name --- flair/__init__.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/flair/__init__.py b/flair/__init__.py index 6dfdead37..ac9770d34 100644 --- a/flair/__init__.py +++ b/flair/__init__.py @@ -24,12 +24,12 @@ """ # Get the device from the environment variable -flair_device = os.environ.get("FLAIR_DEVICE") +device_id = os.environ.get("FLAIR_DEVICE") # global variable: device -if torch.cuda.is_available() and flair_device != "cpu": +if torch.cuda.is_available() and device_id != "cpu": # No need for correctness checks, torch is doing it - device = torch.device(f"cuda:{flair_device}") if flair_device else torch.device("cuda:0") + device = torch.device(f"cuda:{device_id}") if device_id else torch.device("cuda:0") else: device = torch.device("cpu")