Addresssing comments

ml6team · Mar 12, 2024 · 2d87d8c · 2d87d8c
1 parent b00a615
commit 2d87d8c
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/fondant/components/retrieve_from_faiss_by_prompt/src/main.py b/src/fondant/components/retrieve_from_faiss_by_prompt/src/main.py
@@ -48,12 +48,13 @@ def __init__(  # PLR0913
 
     def setup(self) -> Client:
         """Setup LocalCudaCluster if gpu is available."""
+        dask.config.set({"dataframe.convert-string": False})
+        dask.config.set({"distributed.worker.daemon": False})
+
         if self.device == "cuda":
             cluster = LocalCUDACluster()
             return Client(cluster)
 
-        dask.config.set({"dataframe.convert-string": False})
-        dask.config.set({"distributed.worker.daemon": False})
         total_memory = (os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")) / (
             1024**3
         )