diff --git a/tests/test_models.py b/tests/test_models.py
index 0715ebf3..3297ac09 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -15,91 +15,9 @@ torch.manual_seed(SEED)
 np.random.seed(SEED)
 
 
-def get_model():
-    # set up model with these hyperparams
-    root_dir = os.path.dirname(os.path.realpath(__file__))
-    hparams = Namespace(**{'drop_prob': 0.2,
-                           'batch_size': 32,
-                           'in_features': 28*28,
-                           'learning_rate': 0.001*8,
-                           'optimizer_name': 'adam',
-                           'data_root': os.path.join(root_dir, 'mnist'),
-                           'out_features': 10,
-                           'hidden_dim': 1000})
-    model = LightningTemplateModel(hparams)
-
-    return model, hparams
-
-
-def get_exp(debug=True):
-    # set up exp object without actually saving logs
-    root_dir = os.path.dirname(os.path.realpath(__file__))
-    exp = Experiment(debug=debug, save_dir=root_dir, name='tests_tt_dir')
-    return exp
-
-
-def init_save_dir():
-    root_dir = os.path.dirname(os.path.realpath(__file__))
-    save_dir = os.path.join(root_dir, 'save_dir')
-
-    if os.path.exists(save_dir):
-        shutil.rmtree(save_dir)
-
-    os.makedirs(save_dir, exist_ok=True)
-
-    return save_dir
-
-
-def clear_save_dir():
-    root_dir = os.path.dirname(os.path.realpath(__file__))
-    save_dir = os.path.join(root_dir, 'save_dir')
-    if os.path.exists(save_dir):
-        shutil.rmtree(save_dir)
-
-
-def load_model(exp, save_dir):
-
-    # load trained model
-    tags_path = exp.get_data_path(exp.name, exp.version)
-    tags_path = os.path.join(tags_path, 'meta_tags.csv')
-
-    checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
-    weights_dir = os.path.join(save_dir, checkpoints[0])
-
-    trained_model = LightningTemplateModel.load_from_metrics(weights_path=weights_dir, tags_csv=tags_path, on_gpu=True)
-
-    assert trained_model is not None, 'loading model failed'
-
-    return trained_model
-
-
-def run_prediction(dataloader, trained_model):
-    # run prediction on 1 batch
-    for batch in dataloader:
-        break
-
-    x, y = batch
-    x = x.view(x.size(0), -1)
-
-    y_hat = trained_model(x)
-
-    # acc
-    labels_hat = torch.argmax(y_hat, dim=1)
-    val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
-    val_acc = torch.tensor(val_acc)
-    val_acc = val_acc.item()
-
-    print(val_acc)
-
-    assert val_acc > 0.70, f'this model is expected to get > 0.7 in test set (it got {val_acc})'
-
-
-def assert_ok_acc(trainer):
-    # this model should get 0.80+ acc
-    acc = trainer.tng_tqdm_dic['val_acc']
-    assert acc > 0.70, f'model failed to get expected 0.80 validation accuracy. Got: {acc}'
-
-
+# -----------------
+# TESTS
+# -----------------
 def test_cpu_model():
     """
     Make sure model trains on CPU
@@ -306,5 +224,93 @@ def test_amp_gpu_ddp():
     clear_save_dir()
 
 
+# -----------------
+# UTILS
+# -----------------
+def get_model():
+    # set up model with these hyperparams
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    hparams = Namespace(**{'drop_prob': 0.2,
+                           'batch_size': 32,
+                           'in_features': 28*28,
+                           'learning_rate': 0.001*8,
+                           'optimizer_name': 'adam',
+                           'data_root': os.path.join(root_dir, 'mnist'),
+                           'out_features': 10,
+                           'hidden_dim': 1000})
+    model = LightningTemplateModel(hparams)
+
+    return model, hparams
+
+
+def get_exp(debug=True):
+    # set up exp object without actually saving logs
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    exp = Experiment(debug=debug, save_dir=root_dir, name='tests_tt_dir')
+    return exp
+
+
+def init_save_dir():
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    save_dir = os.path.join(root_dir, 'save_dir')
+
+    if os.path.exists(save_dir):
+        shutil.rmtree(save_dir)
+
+    os.makedirs(save_dir, exist_ok=True)
+
+    return save_dir
+
+
+def clear_save_dir():
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    save_dir = os.path.join(root_dir, 'save_dir')
+    if os.path.exists(save_dir):
+        shutil.rmtree(save_dir)
+
+
+def load_model(exp, save_dir):
+
+    # load trained model
+    tags_path = exp.get_data_path(exp.name, exp.version)
+    tags_path = os.path.join(tags_path, 'meta_tags.csv')
+
+    checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
+    weights_dir = os.path.join(save_dir, checkpoints[0])
+
+    trained_model = LightningTemplateModel.load_from_metrics(weights_path=weights_dir, tags_csv=tags_path, on_gpu=True)
+
+    assert trained_model is not None, 'loading model failed'
+
+    return trained_model
+
+
+def run_prediction(dataloader, trained_model):
+    # run prediction on 1 batch
+    for batch in dataloader:
+        break
+
+    x, y = batch
+    x = x.view(x.size(0), -1)
+
+    y_hat = trained_model(x)
+
+    # acc
+    labels_hat = torch.argmax(y_hat, dim=1)
+    val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+    val_acc = torch.tensor(val_acc)
+    val_acc = val_acc.item()
+
+    print(val_acc)
+
+    assert val_acc > 0.70, f'this model is expected to get > 0.7 in test set (it got {val_acc})'
+
+
+def assert_ok_acc(trainer):
+    # this model should get 0.80+ acc
+    acc = trainer.tng_tqdm_dic['val_acc']
+    assert acc > 0.70, f'model failed to get expected 0.70 validation accuracy. Got: {acc}'
+
+
 if __name__ == '__main__':
     pytest.main([__file__])