diff --git a/README.md b/README.md
index f768fe9..843a9f5 100644
--- a/README.md
+++ b/README.md
@@ -35,6 +35,8 @@ With `uv` installed,
 uv add "git+https://github.com/MatchCake/MatchCake-Opt"
 ```
 
+To install the package with cu128 (CUDA), add `--extra cu128` to the installation commands above.
+
 
 ### For developers
 
@@ -79,7 +81,7 @@ class LinearNN(ClassificationModel):
         RangeParameterConfig(
             name="n_neurons",
             parameter_type="int",
-            bounds=(4, 16),
+            bounds=(4, 2048),
         ),
     ]
 
@@ -93,7 +95,7 @@ class LinearNN(ClassificationModel):
     ):
         super().__init__(input_shape=input_shape, output_shape=output_shape, learning_rate=learning_rate, **kwargs)
         self.save_hyperparameters("learning_rate", "n_neurons")
-        self.nn = torch.Sequential(
+        self.nn = torch.nn.Sequential(
             torch.nn.Flatten(),
             torch.nn.LazyLinear(n_neurons), 
             torch.nn.ReLU(), 
@@ -107,7 +109,7 @@ class LinearNN(ClassificationModel):
     def output_size(self):
         return int(np.prod(self.output_shape))
 
-datamodule = DataModule.from_dataset_name("Digits2D")
+datamodule = DataModule.from_dataset_name("Digits2D", fold_id=0)
 automl_pipeline = AutoMLPipeline(model_cls=LinearNN, datamodule=datamodule)
 automl_pipeline.run()
 lt_pipeline, metrics = automl_pipeline.run_best_pipeline()
diff --git a/src/matchcake_opt/tr_pipeline/automl_pipeline.py b/src/matchcake_opt/tr_pipeline/automl_pipeline.py
index 09c1832..e933431 100644
--- a/src/matchcake_opt/tr_pipeline/automl_pipeline.py
+++ b/src/matchcake_opt/tr_pipeline/automl_pipeline.py
@@ -19,8 +19,8 @@
 
 class AutoMLPipeline:
     DEFAULT_AUTOML_ITERATIONS = 32
-    DEFAULT_INNER_MAX_EPOCHS = 10_000
-    DEFAULT_INNER_MAX_TIME = "90:00:00:00"
+    DEFAULT_INNER_MAX_EPOCHS = 1024
+    DEFAULT_INNER_MAX_TIME = "00:01:00:00"
     DEFAULT_AUTOML_OVERWRITE_FIT = False
 
     @classmethod