diff --git a/experiments/5_steps/config_16_layer_16_hidden_refined.yaml b/experiments/5_steps/config_16_layer_16_hidden_refined.yaml
index 91072c1..0c74089 100644
--- a/experiments/5_steps/config_16_layer_16_hidden_refined.yaml
+++ b/experiments/5_steps/config_16_layer_16_hidden_refined.yaml
@@ -58,7 +58,7 @@ data:
   class_path: ThermalSolver.graph_datamodule_unsteady.GraphDataModule
   init_args:
     hf_repo: "SISSAmathLab/thermal-conduction-unsteady"
-    split_name: "easy.refined"
+    split_name: "basic.refined"
     n_elements: 100
     batch_size: 32
     train_size: 0.7
diff --git a/experiments/5_steps/config_32_layer_16_hidden_refined.yaml b/experiments/5_steps/config_32_layer_16_hidden_refined.yaml
index 1d53bc2..a9d6ff6 100644
--- a/experiments/5_steps/config_32_layer_16_hidden_refined.yaml
+++ b/experiments/5_steps/config_32_layer_16_hidden_refined.yaml
@@ -39,7 +39,7 @@ trainer:
   min_steps: null
   overfit_batches: 0.0
   log_every_n_steps: 0
-  accumulate_grad_batches: 2
+  accumulate_grad_batches: 1
   default_root_dir: null
   gradient_clip_val: 1.0
 
@@ -58,9 +58,9 @@ data:
   class_path: ThermalSolver.graph_datamodule_unsteady.GraphDataModule
   init_args:
     hf_repo: "SISSAmathLab/thermal-conduction-unsteady"
-    split_name: "easy.refined"
+    split_name: "basic.refined"
     n_elements: 100
-    batch_size: 16
+    batch_size: 24
     train_size: 0.7
     val_size: 0.2
     test_size: 0.1
diff --git a/experiments/5_steps/config_8_layer_16_hidden_refined.yaml b/experiments/5_steps/config_8_layer_16_hidden_refined.yaml
index 29b97c3..3fc10d5 100644
--- a/experiments/5_steps/config_8_layer_16_hidden_refined.yaml
+++ b/experiments/5_steps/config_8_layer_16_hidden_refined.yaml
@@ -24,7 +24,7 @@ trainer:
       init_args:
         monitor: val/loss
         mode: min
-        patience: 20
+        patience: 30
         verbose: false
   max_epochs: 1000
   min_epochs: null
@@ -50,7 +50,7 @@ data:
   class_path: ThermalSolver.graph_datamodule_unsteady.GraphDataModule
   init_args:
     hf_repo: "SISSAmathLab/thermal-conduction-unsteady"
-    split_name: "easy.refined"
+    split_name: "basic.refined"
     n_elements: 100
     batch_size: 32
     train_size: 0.7