diff --git a/__pycache__/bindataset.cpython-38.pyc b/__pycache__/bindataset.cpython-38.pyc
deleted file mode 100644
index b8f372144a8007149ebb43b4f9565c119f7786ea..0000000000000000000000000000000000000000
Binary files a/__pycache__/bindataset.cpython-38.pyc and /dev/null differ
diff --git a/config.yml b/config.yml
index b448853a01b2640dd32850e38bfedcac7db14d04..8537f702ed1ebf05d9dd19a240f2f44bb12588fe 100644
--- a/config.yml
+++ b/config.yml
@@ -28,6 +28,7 @@ LinearRegression:
   # Bias in {True, False}
   Bias: True
 
-
+#Name of directory containing logs
+LogDir: ./logs/
 
 
diff --git a/create_submission.py b/create_submission.py
index 45aebaa5f7c5abf4759c81ba643463e3cdba50f4..126caf8547e2c640ba89640169168c6f1ae8af22 100644
--- a/create_submission.py
+++ b/create_submission.py
@@ -19,11 +19,19 @@ import datetime
 # External imports
 import tqdm
 import torch
+import torch.nn as nn
 
 # Local imports
 import bindataset as dataset
+import dataloader
 
-def create_submission(model, transform):
+def dummy_model(X):
+    # X is a (B, T, N) tensor
+    # As a dummy model, say, we average all the environmental measures
+    # Divided by a magic number
+    return X[:, :, 4:].mean(dim=2) / 26  # This is (B, T)
+
+def create_submission(model, transform, device):
     step_days = 10
     batch_size = 1024
     # We make chunks of num_days consecutive samples; As our dummy predictor
@@ -33,9 +41,6 @@ def create_submission(model, transform):
     num_days = 365
     num_workers = 7
 
-    use_cuda = torch.cuda.is_available()
-    device = torch.device("cuda") if use_cuda else torch.device("cpu")
-
     # Build the dataloaders
     logging.info("Building the dataloader")
 
@@ -63,11 +68,10 @@ def create_submission(model, transform):
         # days of the same location then followed by consecutive days of the
         # next location and so on
         chunk_size = batch_size * num_days
-
         with torch.no_grad():
             for X in tqdm.tqdm(test_loader):
-                X.to(device)
-
+                X = X.to(device)
+                print(X.shape)
                 #############################################
                 # This is where you inject your knowledge
                 # About your model
@@ -137,4 +141,34 @@ def create_submission(model, transform):
 
 if __name__ == "__main__":
     logging.basicConfig(stream=sys.stdout, level=logging.INFO, format="%(message)s")
-    test()
\ No newline at end of file
+    use_cuda = torch.cuda.is_available()
+    if use_cuda :
+        device = torch.device('cuda')
+    else :
+        device = toch.device('cpu')
+
+    model_path = "logs/LinearRegression_5/best_model.pt"
+
+    model = nn.Sequential(
+        nn.Linear(14,8,False),
+        nn.ReLU(),
+        nn.Linear(8, 35, True),
+        nn.ReLU(),
+        nn.Linear(35,35,True),
+        nn.ReLU(),
+        nn.Linear(35,35,True),
+        nn.ReLU(),
+        nn.Linear(35,35,True),
+        nn.ReLU(),
+        nn.Linear(35,35,True),
+        nn.ReLU(),
+        nn.Linear(35,1, True),
+        nn.ReLU()
+    )
+
+    model = model.to(device)
+
+    model.load_state_dict(torch.load(model_path))
+
+    #create_submission(model, dataloader.transform_remove_space_time(), device)
+    create_submission(model, None, device)
\ No newline at end of file
diff --git a/dataloader.py b/dataloader.py
index 6f9f4f6bcc4d4d1929e0ce58e716f246a9fd2e03..0f131eaa66a207f90f12831ea1634a59f767dce7 100644
--- a/dataloader.py
+++ b/dataloader.py
@@ -200,7 +200,7 @@ if __name__ == "__main__":
         num_days = num_days,
         batch_size = batch_size,
         num_workers = num_workers,
-        pin_memory = False,
+        pin_memory = True,
         valid_ratio = valid_ratio,
         overwrite_index=True,
         max_num_samples=max_num_samples,
@@ -241,6 +241,8 @@ if __name__ == "__main__":
         print(check_max(tensor, index))
         print("="*30)
 
+    print(X.shape)
+
     check_info(X,0) #latitude
     check_info(X,1) #longitude
     check_info(X,2) #depth
diff --git a/logs/dataloader_unit_test.log b/logs/dataloader_unit_test.log
deleted file mode 100644
index 950d264681bca10e805befeff6e2726f308fb547..0000000000000000000000000000000000000000
--- a/logs/dataloader_unit_test.log
+++ /dev/null
@@ -1,61 +0,0 @@
-INFO:root:====> Test dataloader
-INFO:root:= Dataloaders
-INFO:root:  - Dataset creation
-INFO:root:The loaded dataset contains 25 latitudes, 37 longitudes, 28 depths and 2222 time points
-INFO:root:Generating the index
-INFO:root:Loading the index from sub_2CMEMS-MEDSEA-2010-2016-training.nc.bin_index.idx
-INFO:root:  - Loaded a dataset with 50154984 samples
-INFO:root:  - Splitting the data in training and validation sets
-INFO:root:Generating the subset files from 50154984 samples
-INFO:root:  - Subset dataset
-INFO:root:The loaded dataset contains 25 latitudes, 37 longitudes, 28 depths and 2222 time points
-INFO:root:Loading the index from sub_2CMEMS-MEDSEA-2010-2016-training.nc.bin_index.idx
-INFO:root:The loaded dataset contains 25 latitudes, 37 longitudes, 28 depths and 2222 time points
-INFO:root:Loading the index from sub_2CMEMS-MEDSEA-2010-2016-training.nc.bin_index.idx
-INFO:root:  - The train fold has 40123681 samples
-INFO:root:  - The valid fold has 10031303 samples
-INFO:root:Got a minibatch of size torch.Size([128, 1, 18]) -> torch.Size([128, 1])
-INFO:root:====> Test dataloader
-INFO:root:= Dataloaders
-INFO:root:  - Dataset creation
-INFO:root:The loaded dataset contains 25 latitudes, 37 longitudes, 28 depths and 2222 time points
-INFO:root:Generating the index
-INFO:root:Loading the index from sub_2CMEMS-MEDSEA-2010-2016-training.nc.bin_index.idx
-INFO:root:  - Loaded a dataset with 50154984 samples
-INFO:root:  - Splitting the data in training and validation sets
-INFO:root:Generating the subset files from 50154984 samples
-INFO:root:  - Subset dataset
-INFO:root:The loaded dataset contains 25 latitudes, 37 longitudes, 28 depths and 2222 time points
-INFO:root:Loading the index from sub_2CMEMS-MEDSEA-2010-2016-training.nc.bin_index.idx
-INFO:root:The loaded dataset contains 25 latitudes, 37 longitudes, 28 depths and 2222 time points
-INFO:root:Loading the index from sub_2CMEMS-MEDSEA-2010-2016-training.nc.bin_index.idx
-INFO:root:  - The train fold has 40125966 samples
-INFO:root:  - The valid fold has 10029018 samples
-INFO:root:Got a minibatch of size torch.Size([128, 1, 18]) -> torch.Size([128, 1])
-INFO:root:====> Test dataloader
-INFO:root:= Dataloaders
-INFO:root:  - Dataset creation
-INFO:root:The loaded dataset contains 25 latitudes, 37 longitudes, 28 depths and 2222 time points
-INFO:root:Generating the index
-INFO:root:====> Test dataloader
-INFO:root:= Dataloaders
-INFO:root:  - Dataset creation
-INFO:root:The loaded dataset contains 25 latitudes, 37 longitudes, 28 depths and 2222 time points
-INFO:root:Generating the index
-INFO:root:====> Test dataloader
-INFO:root:= Dataloaders
-INFO:root:  - Dataset creation
-INFO:root:The loaded dataset contains 25 latitudes, 37 longitudes, 28 depths and 2222 time points
-INFO:root:Generating the index
-INFO:root:Loading the index from sub_2CMEMS-MEDSEA-2010-2016-training.nc.bin_index.idx
-INFO:root:  - Loaded a dataset with 50154984 samples
-INFO:root:  - Splitting the data in training and validation sets
-INFO:root:Generating the subset files from 50154984 samples
-INFO:root:  - Subset dataset
-INFO:root:The loaded dataset contains 25 latitudes, 37 longitudes, 28 depths and 2222 time points
-INFO:root:Loading the index from sub_2CMEMS-MEDSEA-2010-2016-training.nc.bin_index.idx
-INFO:root:The loaded dataset contains 25 latitudes, 37 longitudes, 28 depths and 2222 time points
-INFO:root:Loading the index from sub_2CMEMS-MEDSEA-2010-2016-training.nc.bin_index.idx
-INFO:root:  - The train fold has 40126890 samples
-INFO:root:  - The valid fold has 10028094 samples
-INFO:root:Got a minibatch of size torch.Size([128, 1, 18]) -> torch.Size([128, 1])
diff --git a/main.py b/main.py
index a29d8dc042daad917efac0f5f85c721257e368e4..a759e7eabcd93d7a51e0e420752854630a42b387 100644
--- a/main.py
+++ b/main.py
@@ -1,18 +1,23 @@
+#Internal imports
 import dataloader
 import model
 import test
 from train import train
-import yaml
 import losses
 import optimizers
+import create_submission
+import utils
+
+#External imports
+import yaml
 import torch
 import logging
 import torch.optim
 import torch.nn as nn
-import create_submission
+import os
 
-def optimizer(cfg, model):
-    result = {"Adam" : torch.optim.Adam(model.parameters())}
+def optimizer(cfg, network):
+    result = {"Adam" : torch.optim.Adam(network.parameters())}
     return result[cfg["Optimizer"]]
 
 if __name__ == "__main__":
@@ -47,39 +52,57 @@ if __name__ == "__main__":
     else :
         device = toch.device('cpu')
 
-    #model = model.build_model(cfg, 18)
+    #network = network.build_network(cfg, 18)
 
-    model = nn.Sequential(
+    network = nn.Sequential(
         nn.Linear(14,8,False),
         nn.ReLU(),
-        nn.Linear(8, 8, True),
+        nn.Linear(8, 35, True),
+        nn.ReLU(),
+        nn.Linear(35,35,True),
+        nn.ReLU(),
+        nn.Linear(35,35,True),
+        nn.ReLU(),
+        nn.Linear(35,35,True),
         nn.ReLU(),
-        nn.Linear(8,35,True),
+        nn.Linear(35,35,True),
         nn.ReLU(),
-        nn.Linear(35,1, True)
+        nn.Linear(35,1, True),
+        nn.ReLU()
     )
-    model = model.to(device)
 
-    for param in list(model.parameters()):
+    def init_xavier(module):
+        if type(module)==nn.Linear:
+            nn.init.xavier_uniform_(module.weight)
+    network = network.to(device)
+
+    """
+    for param in list(network.parameters()):
         param = 1
+    """
 
     f_loss = losses.RMSLELoss()
 
-    optimizer = optimizer(cfg, model)
-    #optimizer = torch.optim.Adam((model.parameters()), lr = 10000)
-  
+    optimizer = optimizer(cfg, network)
+
+    logdir = utils.create_unique_logpath(cfg["LogDir"], cfg["Model"]["Name"])
+    network_checkpoint = model.ModelCheckpoint(logdir + "/best_model.pt", network)
 
     for t in range(cfg["Training"]["Epochs"]):
         torch.autograd.set_detect_anomaly(True)
         print("Epoch {}".format(t))
-        train(model, train_loader, f_loss, optimizer, device)
+        train(network, train_loader, f_loss, optimizer, device)
 
 
-        #print(list(model.parameters())[0].grad)
-        val_loss = test.test(model, valid_loader, f_loss, device)
+        #print(list(network.parameters())[0].grad)
+        val_loss = test.test(network, valid_loader, f_loss, device)
+
+        network_checkpoint.update(val_loss)
+
         print(" Validation : Loss : {:.4f}".format(val_loss))
 
-    create_submission.create_submission(model, None)
+
+    create_submission.create_submission(network, None)
     """
     logdir = generate_unique_logpath(top_logdir, "linear")
     print("Logging to {}".format(logdir))
diff --git a/model.py b/model.py
index 4e6830cb8da450b16918e62505c2cd25403919e0..a40839c7fea2ec3d932515f24fff463ec3f7365a 100644
--- a/model.py
+++ b/model.py
@@ -20,6 +20,17 @@ class LinearRegression(nn.Module):
 def build_model(cfg, input_size):    
     return eval(f"{cfg['Model']['Name']}(cfg, input_size)")
 
+class ModelCheckpoint:
+    def __init__(self, filepath, model):
+        self.min_loss = None
+        self.filepath = filepath
+        self.model = model
+
+    def update(self, loss):
+        if (self.min_loss is None) or (loss < self.min_loss):
+            print("Saving a better model")
+            torch.save(self.model.state_dict(), self.filepath)
+            self.min_loss = loss
 
 if __name__== "__main__":
     import yaml
diff --git a/submission.csv b/submission.csv
new file mode 100644
index 0000000000000000000000000000000000000000..b1f81b8ddbf4ebdd3e2b54e44d602a3ebeb356fc
--- /dev/null
+++ b/submission.csv
@@ -0,0 +1 @@
+Id,Predicted
diff --git a/utils.py b/utils.py
index b6b3e52f630f8b2af68e562cffc74fa31058cbfc..0c585e715e1357f627c86f3a759f9c131fa34768 100644
--- a/utils.py
+++ b/utils.py
@@ -9,3 +9,9 @@ def generate_unique_logpath(logdir, raw_run_name):
             return log_path
         i = i + 1
 
+def create_unique_logpath(top_logdir, raw_run_name):
+    if not os.path.exists(top_logdir):
+        os.mkdir(top_logdir)
+    logdir = generate_unique_logpath(top_logdir, raw_run_name)
+    os.mkdir(logdir)
+    return logdir
\ No newline at end of file