fix: prefer device_id name

flairNLP · Mar 8, 2024 · cc1d938 · cc1d938
1 parent f074b98
commit cc1d938
Showing 1 changed file with 3 additions and 3 deletions.
diff --git 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")