diff --git a/experiment1/model.py b/experiment1/model.py
index 7ca00fc..79d3f90 100644
--- a/experiment1/model.py
+++ b/experiment1/model.py
@@ -35,8 +35,8 @@ def forward(self, x):
 
     def step(self, batch, batch_idx, *args, **kwargs) -> T.Tensor:
         # Unpacking
-        samples = batch["samples"]
-        targets = batch["targets"]
+        samples = batch["sample"]
+        targets = batch["target"]
 
         # Forward
         targets_pred = self(samples)
diff --git a/experiment1/trainer.py b/experiment1/trainer.py
index f7fded5..df9cda8 100644
--- a/experiment1/trainer.py
+++ b/experiment1/trainer.py
@@ -4,12 +4,12 @@
 from model import Model1
 
 # Data loaders
-dl_train = DataLoader(NumberAdd(20000), batch_size=8, shuffle=True, num_workers=1)
-dl_valid = DataLoader(NumberAdd(1000), batch_size=8, shuffle=True, num_workers=1)
-dl_tests = DataLoader(NumberAdd(500), batch_size=8, shuffle=True, num_workers=1)
+dl_train = DataLoader(NumberAdd(20000), batch_size=8, shuffle=True, num_workers=0)
+dl_valid = DataLoader(NumberAdd(1000), batch_size=8, shuffle=True, num_workers=0)
+dl_tests = DataLoader(NumberAdd(500), batch_size=8, shuffle=True, num_workers=0)
 
 # Model
-model = Model1()
+model = Model1(lr=1e-4, optim_conf={"momentum": 0.9})
 
 # Trainer
 trainer = Trainer(max_epochs=500, progress_bar_refresh_rate=20)