diff --git a/experiment1/model.py b/experiment1/model.py
index eb49972..5d78b6c 100644
--- a/experiment1/model.py
+++ b/experiment1/model.py
@@ -30,7 +30,8 @@ def __init__(self, lr, optim_conf: dict={}):
     def forward(self, x):
         y0 = self.adder(x)
         y1 = self.substracter(x)
-        y = self.weights_adder * y0 + self.weights_substracter * y1
+        y = (T.sigmoid(self.weights_adder) * y0 + 
+             T.sigmoid(self.weights_substracter) * y1)
         return y
 
     def step(self, batch, batch_idx, *args, **kwargs) -> T.Tensor:
@@ -59,5 +60,7 @@ def configure_optimizers(self):
         return optim.SGD(self.parameters(), lr=self.lr, **self.optim_conf)
 
     def on_epoch_end(self) -> None:
-        self.log("weights/weights_adder", self.weights_adder)
-        self.log("weights/weights_substractor", self.weights_substracter)
\ No newline at end of file
+        self.log("weights/wa", self.weights_adder)
+        self.log("weights/ws", self.weights_substracter)
+        self.log("weights/sigma(wa)", T.sigmoid(self.weights_adder))
+        self.log("weights/sigma(ws)", T.sigmoid(self.weights_substracter))
\ No newline at end of file
diff --git a/experiment1/trainer.py b/experiment1/trainer.py
index faca364..c247bf2 100644
--- a/experiment1/trainer.py
+++ b/experiment1/trainer.py
@@ -1,5 +1,7 @@
 from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import ModelCheckpoint
 from torch.utils.data.dataloader import DataLoader
+
 from dataset import NumberAdd
 from model import Model1
 
@@ -12,6 +14,6 @@
 model = Model1(lr=1e-4, optim_conf={"momentum": 0.9})
 
 # Trainer
-trainer = Trainer(max_epochs=500)
+trainer = Trainer(max_epochs=500, callbacks=[ModelCheckpoint(monitor="loss/valid")])
 trainer.fit(model=model, train_dataloaders=dl_train, val_dataloaders=dl_valid)