Enable GPU usage

src/pipeline.ipynb

@@ -342,6 +342,10 @@
     "        \"faiss_index_path\":\"hf://datasets/fondant-ai/datacomp-small-clip/faiss\",\n",
     "        \"num_images\": 2\n",
     "    },\n",
+    "    resources=Resources(\n",
+    "        accelerator_number=number_of_accelerators,\n",
+    "        accelerator_name=accelerator_name,\n",
+    "    ),\n",
     ")\n",
     "\n",
     "images = image_urls.apply(\n",

src/pipeline.py

@@ -1,12 +1,11 @@
 """Pipeline used to create a stable diffusion dataset from a set of initial prompts."""
 from pathlib import Path
-
+import torch
 import pyarrow as pa
-from fondant.pipeline import Pipeline
+from fondant.pipeline import Pipeline, Resources
 
 from components.generate_prompts import GeneratePromptsComponent
 
-
 BASE_PATH = "./data_dir"
 # Create data directory if it doesn't exist
 Path(BASE_PATH).mkdir(parents=True, exist_ok=True)
@@ -21,7 +20,7 @@
 prompts = pipeline.read(
     GeneratePromptsComponent,
     arguments={
-        "n_rows_to_load": 10
+        "n_rows_to_load": 50
     },  # Set to 10 for small scale testing, set to None to load all rows
 )
 
@@ -32,6 +31,9 @@
         "faiss_index_path": "hf://datasets/fondant-ai/datacomp-small-clip/faiss",
         "num_images": 2,
     },
+    resources=Resources(
+        accelerator_name="GPU", accelerator_number=torch.cuda.device_count()
+    ),
 )
 
 images = image_urls.apply(
@@ -54,6 +56,9 @@
         "batch_size": 8,
         "max_new_tokens": 50,
     },
+    resources=Resources(
+        accelerator_name="GPU", accelerator_number=torch.cuda.device_count()
+    ),
 )
 
 segmentations = captions.apply(
@@ -62,6 +67,9 @@
         "model_id": "openmmlab/upernet-convnext-small",
         "batch_size": 8,
     },
+    resources=Resources(
+        accelerator_name="GPU", accelerator_number=torch.cuda.device_count()
+    ),
 )