diff --git a/data/kgdataset.py b/data/kgdataset.py
index d12bb93..c34dac7 100644
--- a/data/kgdataset.py
+++ b/data/kgdataset.py
@@ -46,8 +46,9 @@ def __init__(self, split, transform=None, height=256, width=256, label_csv='trai
         images = np.zeros((num,height,width,3),dtype=np.float32)
         for n in range(num):
             img_file = data_dir + '/{}/'.format(ext) + names[n]
-            print(img_file)
+            #print(img_file)
             jpg_file = img_file.replace('<ext>','jpg')
+            #print(jpg_file)
             image_jpg = cv2.imread(jpg_file,1)
             h,w = image_jpg.shape[0:2]
             if height!=h or width!=w:
diff --git a/split_train.py b/split_train.py
index b7dda18..5fdffff 100644
--- a/split_train.py
+++ b/split_train.py
@@ -1,7 +1,7 @@
 import numpy as np
 import pandas as pds
 
-def split_train_validation(num_val=3000):
+def split_train_validation2(num_val=3000):
     """
     Save train image names and validation image names to csv files
     """
@@ -32,4 +32,8 @@ def split_train_validation(num_val=3000):
     df.to_csv('dataset/validation-%s' % num_val, index=False, header=False)
 
 
-split_train_validation(num_val=3000)
+#split_train_validation2(num_val=3000)
+
+from util import split_train_validation
+
+split_train_validation(num_val = 3000)
diff --git a/trainers/baseline_trainer.py b/trainers/baseline_trainer.py
index 5986305..cceece6 100644
--- a/trainers/baseline_trainer.py
+++ b/trainers/baseline_trainer.py
@@ -1,3 +1,6 @@
+import matplotlib
+matplotlib.use('Agg') # Must be before importing matplotlib.pyplot or pylab! Not require X-server to be running
+
 import torch.nn as nn
 from torch.nn import functional as F
 from torch import optim
@@ -46,8 +49,8 @@
           ]
 batch_size = [
                 # 128, 128,
-                64, 64,
-                40, 40,
+                16, 16,
+                16, 16,
                 # 50
             ]
 
@@ -116,7 +119,7 @@ def load_net(net, name):
 
 def train_baselines():
 
-    train_data, val_data = get_dataloader(96)
+    train_data, val_data = get_dataloader(32)
 
     for model, batch in zip(models, batch_size):
         name = str(model).split()[1]
@@ -127,7 +130,7 @@ def train_baselines():
         # load pre-trained model on train-37479
         net = model(pretrained=True)
         net = nn.DataParallel(net.cuda())
-        load_net(net, name)
+        #load_net(net, name)
         # optimizer = get_optimizer(net, lr=.001, pretrained=True, resnet=True if 'resnet' in name else False)
         optimizer = optim.SGD(lr=.005, momentum=0.9, params=net.parameters(), weight_decay=5e-4)
         train_data.batch_size = batch
diff --git a/util.py b/util.py
index 5753eab..ac2b32b 100644
--- a/util.py
+++ b/util.py
@@ -2,6 +2,8 @@
 from torch.autograd import Variable
 from sklearn.metrics import fbeta_score
 from torch.nn import functional as F
+import matplotlib
+matplotlib.use('Agg') # Must be before importing matplotlib.pyplot or pylab! Not require X-server to be running
 from matplotlib import pyplot as plt
 import pandas as pds
 from datasets import *