diff --git a/.gitignore b/.gitignore
index caa3d08..3aad4aa 100644
--- a/.gitignore
+++ b/.gitignore
@@ -44,5 +44,5 @@ venv/
 .cache/
 
 _cache/
-grid_search/
+grid_search*/
 public.spider.surfsara.nl/
diff --git a/neural_networks/pre_processing_for_ml.py b/neural_networks/pre_processing_for_ml.py
index 451a2e2..cb85e89 100644
--- a/neural_networks/pre_processing_for_ml.py
+++ b/neural_networks/pre_processing_for_ml.py
@@ -11,9 +11,6 @@
 import joblib
 from matplotlib.colors import SymLogNorm
 from torch.utils.data import Dataset
-from concurrent.futures import ThreadPoolExecutor, as_completed
-
-cache = joblib.Memory(location="_cache", verbose=0)
 
 
 def get_rms(data: np.ndarray, maskSup=1e-7):
@@ -95,7 +92,7 @@ def transform_data(root_dir, classes=("continue", "stop"), modes=("", "_val")):
     def process_fits(fits_path):
         with fits.open(fits_path) as hdul:
             image_data = hdul[0].data
-
+        # assert image_data.shape[2] == 2048, (image_data.shape, image_data.shape[2], fits_path)
         transformed = normalize_fits(image_data)
 
         np.savez_compressed(
@@ -111,7 +108,6 @@ def process_fits(fits_path):
         for fits_path in (root_dir / (cls + mode)).glob("*.fits")
     )
 
-
 class FitsDataset(Dataset):
     def __init__(self, root_dir, mode="train"):
         """
@@ -130,6 +126,8 @@ def __init__(self, root_dir, mode="train"):
         ext = ".npz"
         glob_ext = "*" + ext
 
+        self.root_dir = root_dir
+
         for folder in (
             root_dir / (cls + ("" if mode == "train" else "_val")) for cls in classes
         ):
@@ -170,16 +168,18 @@ def __init__(self, root_dir, mode="train"):
         # print(f'{mode}: using the following sources: {sources}')
 
     def compute_statistics(self, normalize):
-        self.mean, self.std = FitsDataset._compute_statistics(self, normalize)
+        cache = Memory(location=self.root_dir / '_cache')
+        cached_compute = cache.cache(FitsDataset._compute_statistics)
+        self.mean, self.std = cached_compute(self, normalize)
         return self.mean, self.std
 
     @staticmethod
-    @cache.cache()
     def _compute_statistics(loader, normalize, verbose=True):
-        if not normalize:
-            return torch.asarray([0]), torch.asarray([1])
         if verbose:
             print("Computing dataset statistics")
+        if not normalize:
+            return torch.asarray([0]), torch.asarray([1])
+
         means = []
         sums_of_squares = []
         f = (lambda x: torch.log(x + 1e-10)) if normalize == 2 else lambda x: x
@@ -281,7 +281,7 @@ def make_histogram(root_dir):
 
 
 if __name__ == "__main__":
-    root = f"public.spider.surfsara.nl/project/lofarvwf/jdejong/CORTEX/calibrator_selection_robertjan/cnn_data"
+    root = f"/scratch-shared/CORTEX/public.spider.surfsara.nl/lofarvwf/jdejong/CORTEX/calibrator_selection_robertjan/cnn_data"
     transform_data(root)
 
     # make_histogram(root)
diff --git a/neural_networks/train_nn.py b/neural_networks/train_nn.py
index 0843800..70278a6 100644
--- a/neural_networks/train_nn.py
+++ b/neural_networks/train_nn.py
@@ -24,8 +24,6 @@
 PROFILE = False
 SEED = None
 
-cache = joblib.Memory(location="_cache", verbose=0)
-
 
 def init_vit(model_name):
     assert model_name == "vit_l_16"
@@ -458,8 +456,7 @@ def main(
         train_dataloader=train_dataloader,
         optimizer=optimizer,
         logging_interval=logging_interval,
-        label_smoothing=label_smoothing,
-        stochastic_smoothing=stochastic_smoothing,
+        smoothing_fn=partial(label_smoother, stochastic=stochastic_smoothing, smoothing_factor=label_smoothing),
     )
     val_step_f = partial(val_step_f, val_dataloader=val_dataloader)
 
@@ -558,6 +555,13 @@ def val_step(model, val_dataloader, global_step, metrics_logger, prepare_data_f)
     return mean_loss, logits, targets
 
 
+def label_smoother(labels: torch.tensor, smoothing_factor: float = 0.1, stochastic: bool = True):
+    smoothing_factor = smoothing_factor - (
+    torch.rand_like(labels) * smoothing_factor * stochastic
+    )
+    smoothed_label = (1 - smoothing_factor) * labels + 0.5 * smoothing_factor
+    return smoothed_label
+
 def train_step(
     model,
     optimizer,
@@ -566,8 +570,7 @@ def train_step(
     global_step,
     logging_interval,
     metrics_logger,
-    label_smoothing=0,
-    stochastic_smoothing=False,
+    smoothing_fn,
 ):
     # print("training")
     model.train()
@@ -578,11 +581,7 @@ def train_step(
         global_step += 1
 
         data, labels = prepare_data_f(data, labels)
-        # Stochastic smoothing factor
-        smoothing_factor = label_smoothing - (
-            torch.rand_like(labels) * label_smoothing * stochastic_smoothing
-        )
-        smoothed_label = (1 - smoothing_factor) * labels + 0.5 * smoothing_factor
+        smoothed_label = smoothing_fn(labels)
         data = augmentation(data)
 
         optimizer.zero_grad(set_to_none=True)