diff --git a/.gitignore b/.gitignore
index 9461af3..cb26ee4 100644
--- a/.gitignore
+++ b/.gitignore
@@ -209,3 +209,4 @@ __marimo__/
 *memray*
 rmm_log.txt
 /nvidia/
+nvidia/
diff --git a/src/segger/cli/segment.py b/src/segger/cli/segment.py
index c5987c1..39e8a5c 100644
--- a/src/segger/cli/segment.py
+++ b/src/segger/cli/segment.py
@@ -415,6 +415,7 @@ def segment(
         max_epochs=n_epochs,
         reload_dataloaders_every_n_epochs=1,
         callbacks=callbacks,
+        devices=1,
     )
 
     # Training
diff --git a/src/segger/debug/prediction.py b/src/segger/debug/prediction.py
index 94344f3..8a6633d 100644
--- a/src/segger/debug/prediction.py
+++ b/src/segger/debug/prediction.py
@@ -88,7 +88,7 @@ def run_prediction_only(
     csvlogger = CSVLogger(path_outputs)
     writer = ISTSegmentationWriter(path_outputs, debug=True)
 
-    trainer = Trainer(logger=csvlogger, reload_dataloaders_every_n_epochs=1, callbacks=[writer],)
+    trainer = Trainer(logger=csvlogger, reload_dataloaders_every_n_epochs=1, callbacks=[writer], devices=1,)
     datamodule = ISTDataModule.load_from_checkpoint(path_checkpoint)
     model = LitISTEncoder.load_from_checkpoint(path_checkpoint)