🚧 fix(wip): cuda check

Author: mxchinegod

magnet/ize/memory.py

@@ -2,6 +2,7 @@
 from magnet.utils.globals import _f
 from magnet.utils.milvus import *
 from magnet.utils.data_classes import EmbeddingPayload
+from magnet.utils.globals import Utils
 
 class Embedder:
     """
@@ -24,7 +25,7 @@ class Embedder:
 
     def __init__(self, config, create=False, initialize=False):
         self.config = config
-        self.model = SentenceTransformer(self.config['MODEL'])
+        self.model = SentenceTransformer(self.config['MODEL'], device=Utils.check_cuda())
         _f('info', f'loading into {self.model.device}')
         self.db = MilvusDB(self.config)
         self.db.on()

magnet/utils/globals.py

@@ -154,9 +154,10 @@ def check_cuda(self):
             _f(
                 "info", f"GPU Name - {torch.cuda.get_device_name(0)}"
             )  # 0 is the GPU index
-            return True
+            return 'cuda'
         else:
             _f("warn", "CUDA is not available on this machine.")
+            return 'cpu'
 
     def normalize_text(self, _):
         """