Merge pull request #3350 from flairNLP/device_check

Add a convenience conversion for flair.device
flairNLP · Oct 24, 2023 · a0e5444 · a0e5444
2 parents ed53c42 + c9f7f70
commit a0e5444
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/flair/trainers/trainer.py b/flair/trainers/trainer.py
@@ -486,6 +486,10 @@ def train_custom(
         if epoch == 0:
             self.check_for_and_delete_previous_best_models(base_path)
 
+        # Sanity conversion: if flair.device was set as a string, convert to torch.device
+        if isinstance(flair.device, str):
+            flair.device = torch.device(flair.device)
+
         # -- AmpPlugin -> wraps with AMP
         # -- AnnealingPlugin -> initialize schedulers (requires instantiated optimizer)
         with contextlib.ExitStack() as context_stack: