diff --git a/notebooks/camera_ready/corrupt_data/02 - Finetune Missing Data.ipynb b/notebooks/camera_ready/corrupt_data/02 - Finetune Missing Data.ipynb
new file mode 100644
index 0000000..ea98122
--- /dev/null
+++ b/notebooks/camera_ready/corrupt_data/02 - Finetune Missing Data.ipynb	
@@ -0,0 +1,328 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Missing Data Reconstruction from Pretrained Embeddings\n",
+    "\n",
+    "For this example we're going to build on `01 - Finetune Virtual EVE.ipynb` and create a simpler finetuning set up.\n",
+    "\n",
+    "![Figure 1: Architectural Diagram](assets/architecture_diags_corrupt.svg)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "WARNING: SunpyUserWarning: Importing sunpy.map without its extra dependencies may result in errors.\n",
+      "The following packages are not installed:\n",
+      "['mpl-animators>=1.0.0', 'reproject>=0.9.0']\n",
+      "To install sunpy with these dependencies use `pip install sunpy[map]` or `pip install sunpy[all]` for all extras. \n",
+      "If you installed sunpy via conda, please report this to the community channel: https://matrix.to/#/#sunpy:openastronomy.org [sunpy.util.sysinfo]\n",
+      "WARNING: SunpyUserWarning: Importing sunpy.visualization without its extra dependencies may result in errors.\n",
+      "The following packages are not installed:\n",
+      "['mpl-animators>=1.0.0']\n",
+      "To install sunpy with these dependencies use `pip install sunpy[visualization]` or `pip install sunpy[all]` for all extras. \n",
+      "If you installed sunpy via conda, please report this to the community channel: https://matrix.to/#/#sunpy:openastronomy.org [sunpy.util.sysinfo]\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "import omegaconf\n",
+    "from sdofm.datasets import SDOMLDataModule\n",
+    "import numpy as np"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cfg = omegaconf.OmegaConf.load(\"finetune_corrupt_data.yml\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[* CACHE SYSTEM *] Found cached index data in /mnt/sdoml/cache/aligndata_AIA_FULL_12min.csv.\n",
+      "[* CACHE SYSTEM *] Found cached normalization data in /mnt/sdoml/cache/normalizations_AIA_FULL_12min.json.\n",
+      "[* CACHE SYSTEM *] Found cached HMI mask data in /mnt/sdoml/cache/hmi_mask_512x512.npy.\n"
+     ]
+    }
+   ],
+   "source": [
+    "data_module = SDOMLDataModule(\n",
+    "    hmi_path=None,\n",
+    "    aia_path=(\n",
+    "        os.path.join(\n",
+    "            cfg.data.sdoml.base_directory,\n",
+    "            cfg.data.sdoml.sub_directory.aia,\n",
+    "        )\n",
+    "        if cfg.data.sdoml.sub_directory.aia\n",
+    "        else None\n",
+    "    ),\n",
+    "    eve_path=None,\n",
+    "    components=cfg.data.sdoml.components,\n",
+    "    wavelengths=cfg.data.sdoml.wavelengths,\n",
+    "    ions=cfg.data.sdoml.ions,\n",
+    "    frequency=cfg.data.sdoml.frequency,\n",
+    "    batch_size=cfg.model.opt.batch_size,\n",
+    "    num_workers=cfg.data.num_workers,\n",
+    "    val_months=cfg.data.month_splits.val,\n",
+    "    test_months=cfg.data.month_splits.test,\n",
+    "    holdout_months=cfg.data.month_splits.holdout,\n",
+    "    cache_dir=os.path.join(\n",
+    "        cfg.data.sdoml.base_directory, cfg.data.sdoml.sub_directory.cache\n",
+    "    ),\n",
+    "    min_date=cfg.data.min_date,\n",
+    "    max_date=cfg.data.max_date,\n",
+    "    num_frames=cfg.data.num_frames,\n",
+    "    drop_frame_dim=cfg.data.drop_frame_dim,\n",
+    ")\n",
+    "data_module.setup()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 32,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from sdofm.models import WrapEncoder, ConvTransformerTokensToEmbeddingNeck\n",
+    "from sdofm.benchmarks import reconstruction as bench_recon\n",
+    "import torch.nn.functional as F\n",
+    "from sdofm.constants import ALL_WAVELENGTHS\n",
+    "from sdofm import BaseModule\n",
+    "\n",
+    "class MissingDataModel(BaseModule):\n",
+    "    def __init__(\n",
+    "            self,\n",
+    "            # Backbone parameters\n",
+    "            img_size: int = 512,\n",
+    "            patch_size: int = 16,\n",
+    "            embed_dim: int = 128,\n",
+    "            num_frames: int = 1,\n",
+    "            # for finetuning\n",
+    "            backbone: object = None,\n",
+    "            freeze_encoder: bool = True,\n",
+    "            # all else\n",
+    "            *args,\n",
+    "            **kwargs,\n",
+    "        ):\n",
+    "            super().__init__(*args, **kwargs)\n",
+    "\n",
+    "            self.backbone = backbone\n",
+    "\n",
+    "            self.masking_ratio = 0.75\n",
+    "            self.validation_metrics = []\n",
+    "\n",
+    "            if freeze_encoder:\n",
+    "                self.backbone.autoencoder.blocks.eval()\n",
+    "                for param in self.backbone.autoencoder.blocks.parameters():\n",
+    "                    param.requires_grad = False\n",
+    "\n",
+    "            self.simulated_corrupt_wavelength = 5\n",
+    "\n",
+    "            # As this is a reconstruction task, something that the MAE\n",
+    "            # was designed to do, we don't require the neck.\n",
+    "            \n",
+    "    def forward_corrupt_data_override(self, imgs, mask_ratio=0.75):\n",
+    "        # corrupt our wavelength by setting it all to 0\n",
+    "        imgs[:,self.simulated_corrupt_wavelength,:,:] = 0\n",
+    "        # continue as normal\n",
+    "        latent, mask, ids_restore = self.backbone.autoencoder.forward_encoder(imgs, mask_ratio)\n",
+    "        pred = self.backbone.autoencoder.forward_decoder(latent, ids_restore)\n",
+    "        loss = self.backbone.autoencoder.forward_loss(imgs, pred, mask)\n",
+    "        return loss, pred, mask\n",
+    "\n",
+    "    def training_step(self, batch, batch_idx):\n",
+    "        # training_step defines the train loop.\n",
+    "        x = batch\n",
+    "        loss, x_hat, mask = self.forward_corrupt_data_override(x, mask_ratio=self.masking_ratio)\n",
+    "        x_hat = self.backbone.autoencoder.unpatchify(x_hat)\n",
+    "        loss = F.mse_loss(x_hat, x)\n",
+    "        self.log(\"train_loss\", loss, on_step=True, on_epoch=True, prog_bar=True)\n",
+    "        return loss\n",
+    "    \n",
+    "    def validation_step(self, batch, batch_idx):\n",
+    "            x = batch\n",
+    "            loss, x_hat, mask = self.backbone.autoencoder(x, mask_ratio=self.masking_ratio)\n",
+    "            x_hat = self.backbone.autoencoder.unpatchify(x_hat)\n",
+    "            loss = F.mse_loss(x_hat, x)\n",
+    "            for i in range(x.shape[0]):\n",
+    "                for frame in range(x.shape[2]):\n",
+    "                    self.validation_metrics.append(\n",
+    "                        bench_recon.get_metrics(\n",
+    "                            x[i, :, frame, :, :], x_hat[i, :, frame, :, :], ALL_WAVELENGTHS\n",
+    "                        )\n",
+    "                    )\n",
+    "\n",
+    "            self.log(\"val_loss\", loss)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 33,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Using <class 'sdofm.datasets.SDOML.SDOMLDataModule'> Data Class\n",
+      "[* CACHE SYSTEM *] Found cached index data in /mnt/sdoml/cache/aligndata_AIA_FULL_12min.csv.\n",
+      "[* CACHE SYSTEM *] Found cached normalization data in /mnt/sdoml/cache/normalizations_AIA_FULL_12min.json.\n",
+      "[* CACHE SYSTEM *] Found cached HMI mask data in /mnt/sdoml/cache/hmi_mask_512x512.npy.\n",
+      "Loading checkpoint...\n",
+      "Done\n"
+     ]
+    }
+   ],
+   "source": [
+    "from pretrain import Pretrainer\n",
+    "MAE = Pretrainer(cfg, logger=None, is_backbone=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 34,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "backbone = MAE.model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 35,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "backbone_params = {}\n",
+    "backbone_params[\"img_size\"] = cfg.model.mae.img_size\n",
+    "backbone_params[\"patch_size\"] = cfg.model.mae.patch_size\n",
+    "backbone_params[\"embed_dim\"] = cfg.model.mae.embed_dim\n",
+    "backbone_params[\"num_frames\"] = cfg.model.mae.num_frames\n",
+    "\n",
+    "model = MissingDataModel(\n",
+    "    # backbone\n",
+    "    **backbone_params,\n",
+    "    # backbone\n",
+    "    backbone=backbone,\n",
+    "    hyperparam_ignore=[\"backbone\"],\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 36,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Trainer will use only 1 of 4 GPUs because it is running inside an interactive / notebook environment. You may try to set `Trainer(devices=4)` but please note that multi-GPU inside interactive / notebook environments is considered experimental and unstable. Your mileage may vary.\n",
+      "GPU available: True (cuda), used: True\n",
+      "TPU available: False, using: 0 TPU cores\n",
+      "HPU available: False, using: 0 HPUs\n",
+      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3]\n",
+      "\n",
+      "  | Name     | Type | Params | Mode \n",
+      "------------------------------------------\n",
+      "0 | backbone | MAE  | 104 M  | train\n",
+      "------------------------------------------\n",
+      "27.8 M    Trainable params\n",
+      "76.7 M    Non-trainable params\n",
+      "104 M     Total params\n",
+      "418.215   Total estimated model params size (MB)\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "462d1016120945a5b14c8a94c8b60a75",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Sanity Checking: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a9dfb00f12274d8c9310c6298906040c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Training: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/opt/conda/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py:54: Detected KeyboardInterrupt, attempting graceful shutdown...\n"
+     ]
+    }
+   ],
+   "source": [
+    "from lightning.pytorch import Trainer \n",
+    "os.environ['PJRT_DEVICE'] = 'GPU'\n",
+    "trainer = Trainer(max_epochs=2, precision=32)\n",
+    "trainer.fit(model=model, datamodule=data_module)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "base",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.14"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/notebooks/camera_ready/corrupt_data/finetune_corrupt_data.yml b/notebooks/camera_ready/corrupt_data/finetune_corrupt_data.yml
new file mode 100644
index 0000000..b658b40
--- /dev/null
+++ b/notebooks/camera_ready/corrupt_data/finetune_corrupt_data.yml
@@ -0,0 +1,159 @@
+# General
+log_level: 'DEBUG'
+experiment:
+  name: null # Name of wandb run, if null, it will be auto generated.
+  project: "sdofm"
+  task: "finetune"
+  model: "virtualeve"
+  resuming: false
+  checkpoint: null              # this is the wandb run_id of the checkpoint to load
+  backbone:
+    checkpoint: "assets/model.ckpt"
+    model: "mae"
+  seed: 0
+  disable_cuda: false
+  wandb:
+    enable: false
+    entity: ""
+    group: ""
+    job_type: "finetune"
+    tags: []
+    notes: ""
+    output_directory: "wandb_output"
+    log_model: "all"             # can be True (final checkpoint), False (no checkpointing), or "all" (for all epoches)
+  gcp_storage:                  # this will checkpoint all epoches and upload them to a GCP bucket, W&B will store references (TODO: perhaps explain this better)
+    enabled: true
+    bucket: "sdofm-checkpoints"
+  fold: null
+  evaluate: false               # skip training and only evaluate (requires checkpoint to be set)
+  device: null                  # this is set automatically using the disable_cuda flag and torch.cuda.is_available()
+  precision: 'bf16-true'        # (32, 64) for cuda, ('32-true', '16-true', 'bf16-true') for tpu
+  log_n_batches: 1000           # log every n training batches
+  save_results: true            # save full results to file and wandb
+  accelerator: "auto"           # options are "auto", "gpu", "tpu", "ipu", or "cpu"
+  profiler: null                # options are 'XLAProfiler' (TPU), 'PyTorchProfiler', warning: PyTorchProfiler only works on cpu/gpu according to docs
+  distributed:
+    enabled: true               # set to true to use more than one device
+    world_size: "auto"          # The "auto" option recognizes the machine you are on, and selects the appropriate number of accelerators.
+    strategy: "ddp_find_unused_parameters_true"
+  log_every_n_steps: 5
+  
+# Dataset Configuration
+data:
+  min_date: '2011-10-01 00:00:00.00'    # minimum is '2010-09-09 00:00:11.08'
+  max_date: '2011-12-31 23:59:59.99'    # maximum is '2023-05-26 06:36:08.072'
+  month_splits:                 # non selected months will form training set 
+    # train: [1,2,3,4,5,6,7,8,9,10]
+    val: [11]
+    test: [12]
+    holdout: []
+  num_workers: 4               # set appropriately for your machine
+  prefetch_factor: 3            # TODO: not implemented, 2 is default
+  num_frames: 1 # WARNING: This is only read for FINETUNING, model num_frames overrides in BACKBONE
+  drop_frame_dim: false         # Requires num_frames=1, for backwards compatibility
+  sdoml:
+    base_directory: "/mnt/sdoml"
+    sub_directory:
+      hmi: "HMI.zarr"
+      aia: "AIA.zarr"
+      eve: "EVE_legacy.zarr"
+      cache: "cache"
+    components: null            # null for select all magnetic components ["Bx", "By", "Bz"]
+    wavelengths: null           # null for select all wavelengths channels ["131A","1600A","1700A","171A","193A","211A","304A","335A","94A"]
+    ions: null                  # null to select all ion channels ["C III", "Fe IX", "Fe VIII", "Fe X", "Fe XI", "Fe XII", "Fe XIII", "Fe XIV", "Fe XIX", "Fe XV", "Fe XVI", "Fe XVIII", "Fe XVI_2", "Fe XX", "Fe XX_2", "Fe XX_3", "H I", "H I_2", "H I_3", "He I", "He II", "He II_2", "He I_2", "Mg IX", "Mg X", "Mg X_2", "Ne VII", "Ne VIII", "O II", "O III", "O III_2", "O II_2", "O IV", "O IV_2", "O V", "O VI", "S XIV", "Si XII", "Si XII_2"]
+    frequency: '12min'          # smallest is 12min
+    mask_with_hmi_threshold: null # None/null for no mask, float for threshold
+    feature_engineering:
+      enabled: false
+      # dclass: 'HelioProjected'
+
+# model configurations
+model:
+  # PRETRAINERS
+  mae:
+    img_size: 512
+    patch_size: 16
+    num_frames: 1
+    tubelet_size: 1
+    in_chans: 9
+    embed_dim: 512
+    depth: 24
+    num_heads: 16
+    decoder_embed_dim: 512
+    decoder_depth: 8
+    decoder_num_heads: 16
+    mlp_ratio: 4.0
+    norm_layer: 'LayerNorm'
+    norm_pix_loss: False
+    masking_ratio: 0.5
+  samae:
+    # uses all parameters as in mae plus these
+    masking_type: "random"  # 'random' or 'solar_aware'
+    active_region_mu_degs: 15.73
+    active_region_std_degs: 6.14
+    active_region_scale: 1.0
+    active_region_abs_lon_max_degs: 60
+    active_region_abs_lat_max_degs: 60
+  nvae:
+    use_se: true
+    res_dist: true
+    num_x_bits: 8
+    num_latent_scales: 3  # 5
+    num_groups_per_scale: 1  # 16
+    num_latent_per_group: 1  # 10
+    ada_groups: true
+    min_groups_per_scale: 1
+    num_channels_enc: 30
+    num_channels_dec: 30
+    num_preprocess_blocks: 2   # 1
+    num_preprocess_cells: 2
+    num_cell_per_cond_enc: 2
+    num_postprocess_blocks: 2  # 1
+    num_postprocess_cells: 2
+    num_cell_per_cond_dec: 2
+    num_mixture_dec: 1
+    num_nf: 2
+    kl_anneal_portion: 0.3
+    kl_const_portion: 0.0001
+    kl_const_coeff: 0.0001
+    # learning_rate: 1e-2
+    # weight_decay: 3e-4
+    weight_decay_norm_anneal: true
+    weight_decay_norm_init: 1.
+    weight_decay_norm: 1e-2        
+    
+  # FINE-TUNERS
+  autocalibration:
+    freeze_encoder: true
+    num_neck_filters: 32
+    output_dim: 1 # not sure why this is implemented for autocorrelation, should be a scalar
+    loss: "mse" # options: "mse", "heteroscedastic"
+  virtualeve:
+    freeze_encoder: true
+    num_neck_filters: 32
+    cnn_model: "efficientnet_b3"
+    lr_linear: 0.01
+    lr_cnn: 0.0001
+    cnn_dp: 0.75
+    epochs_linear: 2
+    
+  # ML optimization arguments:
+  opt:
+    loss: "mse" # options: "mae", "mse", "mape"
+    scheduler: "constant" #other options: "cosine", "plateau", "exp"
+    scheduler_warmup: 0
+    batch_size: 1
+    learning_rate: 0.0001
+    weight_decay: 3e-4 # 0.0
+    optimiser: "adam"
+    epochs: 5
+    patience: 2
+
+# hydra configuration
+hydra:
+  mode: RUN
+  # run:
+  #   dir: ${data.output_directory}/${now:%Y-%m-%d-%H-%M-%S}
+  # sweep:
+  #   dir: ${hydra.run.dir}
+  #   subdir: ${hydra.job.num}
\ No newline at end of file
diff --git a/notebooks/camera_ready/corrupt_data/pretrain.py b/notebooks/camera_ready/corrupt_data/pretrain.py
new file mode 100644
index 0000000..b8f4eca
--- /dev/null
+++ b/notebooks/camera_ready/corrupt_data/pretrain.py
@@ -0,0 +1,261 @@
+# Main pretraining and evaluation script for SDO-FM
+
+import glob
+import os
+from pathlib import Path
+
+import lightning.pytorch as pl
+import torch
+from lightning.fabric.strategies import XLAFSDPStrategy
+
+import wandb
+from sdofm import utils
+from sdofm.datasets import (
+    SDOMLDataModule,
+    BrightSpotsSDOMLDataModule,
+    HelioProjectedSDOMLDataModule,
+)
+from sdofm.pretraining import MAE, NVAE, SAMAE, BrightSpots
+
+
+class Pretrainer(object):
+    def __init__(self, cfg, logger=None, profiler=None, is_backbone=False):
+        self.cfg = cfg
+        self.logger = logger  # would be wandb but broken
+        self.profiler = profiler  # if profiler is not None else Profiler()
+        self.data_module = None
+        self.model = None
+        self.model_class = None
+        self.data_module_class = SDOMLDataModule
+
+        model_name = (
+            cfg.experiment.model if not is_backbone else cfg.experiment.backbone.model
+        )
+
+        if cfg.data.sdoml.feature_engineering.enabled:
+            match cfg.data.sdoml.feature_engineering.dclass:
+                case "BrightSpots":
+                    self.data_module_class = BrightSpotsSDOMLDataModule
+                case "HelioProjected":
+                    self.data_module_class = HelioProjectedSDOMLDataModule
+        print(f"Using {self.data_module_class} Data Class")
+
+        match model_name:
+            case "mae":
+                self.model_class = MAE
+
+                self.data_module = SDOMLDataModule(
+                    # hmi_path=os.path.join(
+                    #     self.cfg.data.sdoml.base_directory, self.cfg.data.sdoml.sub_directory.hmi
+                    # ),
+                    hmi_path=None,
+                    aia_path=os.path.join(
+                        self.cfg.data.sdoml.base_directory,
+                        self.cfg.data.sdoml.sub_directory.aia,
+                    ),
+                    eve_path=None,
+                    components=self.cfg.data.sdoml.components,
+                    wavelengths=self.cfg.data.sdoml.wavelengths,
+                    ions=self.cfg.data.sdoml.ions,
+                    frequency=self.cfg.data.sdoml.frequency,
+                    batch_size=self.cfg.model.opt.batch_size,
+                    num_workers=self.cfg.data.num_workers,
+                    val_months=self.cfg.data.month_splits.val,
+                    test_months=self.cfg.data.month_splits.test,
+                    holdout_months=self.cfg.data.month_splits.holdout,
+                    cache_dir=os.path.join(
+                        self.cfg.data.sdoml.base_directory,
+                        self.cfg.data.sdoml.sub_directory.cache,
+                    ),
+                    min_date=cfg.data.min_date,
+                    max_date=cfg.data.max_date,
+                    num_frames=cfg.model.mae.num_frames,
+                    drop_frame_dim=cfg.data.drop_frame_dim,
+                )
+                self.data_module.setup()
+
+                if cfg.experiment.resuming or is_backbone:
+                    self.model = self.load_checkpoint(
+                        cfg.experiment.checkpoint
+                        if not is_backbone
+                        else cfg.experiment.backbone.checkpoint
+                    )
+                else:
+                    self.model = self.model_class(
+                        **cfg.model.mae,
+                        optimiser=cfg.model.opt.optimiser,
+                        lr=cfg.model.opt.learning_rate,
+                        weight_decay=cfg.model.opt.weight_decay,
+                    )
+            case "samae":
+                self.model_class = SAMAE
+                self.data_module = self.data_module_class(
+                    hmi_path=None,
+                    aia_path=os.path.join(
+                        cfg.data.sdoml.base_directory, cfg.data.sdoml.sub_directory.aia
+                    ),
+                    eve_path=None,
+                    components=cfg.data.sdoml.components,
+                    wavelengths=cfg.data.sdoml.wavelengths,
+                    ions=cfg.data.sdoml.ions,
+                    frequency=cfg.data.sdoml.frequency,
+                    batch_size=cfg.model.opt.batch_size,
+                    num_workers=cfg.data.num_workers,
+                    val_months=cfg.data.month_splits.val,
+                    test_months=cfg.data.month_splits.test,
+                    holdout_months=cfg.data.month_splits.holdout,
+                    cache_dir=os.path.join(
+                        cfg.data.sdoml.base_directory,
+                        cfg.data.sdoml.sub_directory.cache,
+                    ),
+                    min_date=cfg.data.min_date,
+                    max_date=cfg.data.max_date,
+                    num_frames=cfg.model.mae.num_frames,
+                    drop_frame_dim=cfg.data.drop_frame_dim,
+                )
+                self.data_module.setup()
+
+                if cfg.experiment.resuming or is_backbone:
+                    self.model = self.load_checkpoint(
+                        cfg.experiment.checkpoint
+                        if not is_backbone
+                        else cfg.experiment.backbone.checkpoint
+                    )
+                else:
+                    self.model = self.model_class(
+                        **cfg.model.mae,
+                        **cfg.model.samae,
+                        optimiser=cfg.model.opt.optimiser,
+                        lr=cfg.model.opt.learning_rate,
+                        weight_decay=cfg.model.opt.weight_decay,
+                    )
+
+            case "brightspots":
+                self.model_class = BrightSpots
+                self.data_module = BrightSpotsSDOMLDataModule(
+                    hmi_path=os.path.join(
+                        cfg.data.sdoml.base_directory, cfg.data.sdoml.sub_directory.hmi
+                    ),
+                    aia_path=os.path.join(
+                        cfg.data.sdoml.base_directory, cfg.data.sdoml.sub_directory.aia
+                    ),
+                    eve_path=None,
+                    components=cfg.data.sdoml.components,
+                    wavelengths=cfg.data.sdoml.wavelengths,
+                    ions=cfg.data.sdoml.ions,
+                    frequency=cfg.data.sdoml.frequency,
+                    batch_size=cfg.model.opt.batch_size,
+                    num_workers=cfg.data.num_workers,
+                    blosc_cache=None,  # "/home/walsh/blosc_cache", # disable for TPUs
+                    val_months=cfg.data.month_splits.val,
+                    test_months=cfg.data.month_splits.test,
+                    holdout_months=cfg.data.month_splits.holdout,
+                    cache_dir=os.path.join(
+                        cfg.data.sdoml.base_directory,
+                        cfg.data.sdoml.sub_directory.cache,
+                    ),
+                    min_date=cfg.data.min_date,
+                    max_date=cfg.data.max_date,
+                    num_frames=cfg.data.num_frames,
+                )
+                self.data_module.setup()
+
+                if cfg.experiment.resuming or is_backbone:
+                    self.model = self.load_checkpoint(
+                        cfg.experiment.checkpoint
+                        if not is_backbone
+                        else cfg.experiment.backbone.checkpoint
+                    )
+                else:
+                    self.model = self.model_class(
+                        **cfg.model.brightspots,
+                        optimiser=cfg.model.opt.optimiser,
+                        lr=cfg.model.opt.learning_rate,
+                        weight_decay=cfg.model.opt.weight_decay,
+                    )
+
+            case "nvae":
+                self.model_class = NVAE
+
+                self.data_module = SDOMLDataModule(
+                    hmi_path=os.path.join(
+                        self.cfg.data.sdoml.base_directory,
+                        self.cfg.data.sdoml.sub_directory.hmi,
+                    ),
+                    aia_path=os.path.join(
+                        self.cfg.data.sdoml.base_directory,
+                        self.cfg.data.sdoml.sub_directory.aia,
+                    ),
+                    eve_path=None,
+                    components=self.cfg.data.sdoml.components,
+                    wavelengths=self.cfg.data.sdoml.wavelengths,
+                    ions=self.cfg.data.sdoml.ions,
+                    frequency=self.cfg.data.sdoml.frequency,
+                    batch_size=self.cfg.model.opt.batch_size,
+                    num_workers=self.cfg.data.num_workers,
+                    val_months=self.cfg.data.month_splits.val,
+                    test_months=self.cfg.data.month_splits.test,
+                    holdout_months=self.cfg.data.month_splits.holdout,
+                    cache_dir=os.path.join(
+                        self.cfg.data.sdoml.base_directory,
+                        self.cfg.data.sdoml.sub_directory.cache,
+                    ),
+                    min_date=cfg.data.min_date,
+                    max_date=cfg.data.max_date,
+                    num_frames=cfg.data.num_frames,
+                    drop_frame_dim=cfg.data.num_frames,
+                )
+
+                if cfg.experiment.resuming or is_backbone:
+                    self.model = self.load_checkpoint(
+                        cfg.experiment.checkpoint
+                        if not is_backbone
+                        else cfg.experiment.backbone.checkpoint
+                    )
+                else:
+                    self.model = self.model_class(
+                        **cfg.model.nvae,
+                        optimiser=cfg.model.opt.optimiser,
+                        lr=cfg.model.opt.learning_rate,
+                        weight_decay=cfg.model.opt.weight_decay,
+                        hmi_mask=self.data_module.hmi_mask,
+                    )
+            case _:
+                raise NotImplementedError(
+                    f"Model {cfg.experiment.model} not implemented"
+                )
+
+    def load_checkpoint(self, checkpoint_reference):
+        print("Loading checkpoint...")
+        self.model = self.model_class.load_from_checkpoint(checkpoint_reference)
+        print("Done")
+        return self.model
+
+    def run(self):
+        print("\nPRE-TRAINING\n")
+
+        if self.cfg.experiment.distributed.enabled:
+            trainer = pl.Trainer(
+                devices=self.cfg.experiment.distributed.world_size,
+                accelerator=self.cfg.experiment.accelerator,
+                max_epochs=self.cfg.model.opt.epochs,
+                precision=self.cfg.experiment.precision,
+                profiler=self.profiler,
+                logger=self.logger,
+                enable_checkpointing=True,
+                log_every_n_steps=self.cfg.experiment.log_every_n_steps,
+            )
+        else:
+            trainer = pl.Trainer(
+                accelerator=self.cfg.experiment.accelerator,
+                max_epochs=self.cfg.model.opt.epochs,
+                logger=self.logger,
+            )
+        trainer.fit(model=self.model, datamodule=self.data_module)
+        return trainer
+
+    def evaluate(self):
+        self.trainer.evaluate()
+
+    def test(self):
+        self.trainer.test(ckpt_path="best")