diff --git a/src/train.py b/src/train.py
index de2fd48..4a434ed 100644
--- a/src/train.py
+++ b/src/train.py
@@ -22,7 +22,7 @@ def main():
     wandb.init(
         project="U-Net",
         config=dict(
-            DIR_TRAIN_IMG="/home/lilian/data_disk/lfainsin/smolval2017",
+            DIR_TRAIN_IMG="/home/lilian/data_disk/lfainsin/val2017",
             DIR_VALID_IMG="/home/lilian/data_disk/lfainsin/smoltrain2017/",
             DIR_SPHERE_IMG="/home/lilian/data_disk/lfainsin/spheres/Images/",
             DIR_SPHERE_MASK="/home/lilian/data_disk/lfainsin/spheres/Masks/",
@@ -31,7 +31,7 @@ def main():
             N_CLASSES=1,
             AMP=True,
             PIN_MEMORY=True,
-            BENCHMARK=False,
+            BENCHMARK=True,
             DEVICE="cuda",
             WORKERS=8,
             EPOCHS=5,
diff --git a/src/unet/blocks.py b/src/unet/blocks.py
index b1f7d7f..1f4a854 100644
--- a/src/unet/blocks.py
+++ b/src/unet/blocks.py
@@ -70,10 +70,7 @@ class OutConv(nn.Module):
     def __init__(self, in_channels, out_channels):
         super(OutConv, self).__init__()
 
-        self.conv = nn.Sequential(
-            nn.Conv2d(in_channels, out_channels, kernel_size=1),
-            # nn.Sigmoid(),
-        )
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
 
     def forward(self, x):
         return self.conv(x)