extract: Fix keras.device to call keras.backend.device_scope

plugins/extract/_base.py

@@ -7,9 +7,9 @@
 import typing as T
 from dataclasses import dataclass, field
 
-import keras
 import numpy as np
 import torch
+from keras.backend import device_scope
 
 from lib.logger import parse_class_init
 from lib.multithreading import MultiThread
@@ -424,7 +424,7 @@ def get_device_context(cls, cpu: bool) -> T.ContextManager:
         """
         if cpu:
             logger.debug("CPU mode selected. Returning CPU device context")
-            return keras.device("cpu")
+            return device_scope("cpu")
 
         # TODO apple_silicon + directml
         if get_backend() == "apple_silicon":
@@ -434,10 +434,10 @@ def get_device_context(cls, cpu: bool) -> T.ContextManager:
 
         if torch.cuda.is_available():
             logger.debug("Cuda available. Returning Cuda device context")
-            return keras.device("cuda")
+            return device_scope("cuda")
 
         logger.debug("Cuda not available. Returning CPU device context")
-        return keras.device("cpu")
+        return device_scope("cpu")
 
     # <<< THREADING METHODS >>> #
     def start(self) -> None: