diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 7422e483..9da60704 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -10,6 +10,8 @@ on:
 jobs:
   tests:
     runs-on: ${{ matrix.os }}
+    env:
+      RUN_SLOW_TESTS: 1
     strategy:
       fail-fast: false
       matrix:
@@ -59,14 +61,14 @@ jobs:
       - run: pnpm test:ci
       - run: sh ./scripts/run_tests.sh unzip
 
-      - name: Run simple
-        run: sh ./scripts/run_tests.sh simple
-        if: github.event_name != 'schedule'
-
       - name: Run all
         run: sh ./scripts/run_tests.sh all
         if: github.event_name != 'schedule'
 
+      - name: Run simple
+        run: sh ./scripts/run_tests.sh simple
+        if: github.event_name != 'schedule'
+
       - name: Run launch
         run: sh ./scripts/run_tests.sh launch
         if: github.event_name == 'schedule'
diff --git a/__tests__/vision-classification.spec.js b/__tests__/vision-classification.spec.js
index 13cdfc73..c2c59b2b 100644
--- a/__tests__/vision-classification.spec.js
+++ b/__tests__/vision-classification.spec.js
@@ -43,6 +43,9 @@ test('vision classification simple', async () => {
 test('vision classification all', async () => {
   await page.selectOption('select', 'template-vision-classification')
 
+  await page.check('#include_test-checkbox')
+  expect(await page.isChecked('#include_test-checkbox')).toBeTruthy()
+
   await page.waitForSelector('text=README.md')
   await page.click(':nth-match(:text("Training"), 2)')
 
diff --git a/__tests__/vision-dcgan.spec.js b/__tests__/vision-dcgan.spec.js
index 81753a2e..127458f2 100644
--- a/__tests__/vision-dcgan.spec.js
+++ b/__tests__/vision-dcgan.spec.js
@@ -43,6 +43,9 @@ test('vision dcgan simple', async () => {
 test('vision dcgan all', async () => {
   await page.selectOption('select', 'template-vision-dcgan')
 
+  await page.check('#include_test-checkbox')
+  expect(await page.isChecked('#include_test-checkbox')).toBeTruthy()
+
   await page.waitForSelector('text=README.md')
   await page.click(':nth-match(:text("Training"), 2)')
 
diff --git a/__tests__/vision-segmentation.spec.js b/__tests__/vision-segmentation.spec.js
index 84ab3e10..ef7c2fff 100644
--- a/__tests__/vision-segmentation.spec.js
+++ b/__tests__/vision-segmentation.spec.js
@@ -43,6 +43,9 @@ test('vision segmentation simple', async () => {
 test('vision segmentation all', async () => {
   await page.selectOption('select', 'template-vision-segmentation')
 
+  await page.check('#include_test-checkbox')
+  expect(await page.isChecked('#include_test-checkbox')).toBeTruthy()
+
   await page.waitForSelector('text=README.md')
   await page.click(':nth-match(:text("Training"), 2)')
 
diff --git a/scripts/requirements.txt b/scripts/requirements.txt
index 7a659699..e0c29e52 100644
--- a/scripts/requirements.txt
+++ b/scripts/requirements.txt
@@ -4,3 +4,4 @@ pytorch-ignite>=0.4.2
 pyyaml
 albumentations
 image_dataset_viz
+pytest
diff --git a/scripts/run_tests.sh b/scripts/run_tests.sh
index 84eeee1e..19385ad7 100644
--- a/scripts/run_tests.sh
+++ b/scripts/run_tests.sh
@@ -27,6 +27,7 @@ run_all() {
   for dir in $(find ./dist-tests/*-all -type d)
   do
     cd $dir
+    pytest -vra --color=yes --tb=short test_*.py
     python main.py --data_path ~/data
     cd $CWD
   done
diff --git a/src/components/TabTemplates.vue b/src/components/TabTemplates.vue
index c6f298f8..d0847c83 100644
--- a/src/components/TabTemplates.vue
+++ b/src/components/TabTemplates.vue
@@ -10,16 +10,23 @@
         @change.prevent="downloadTemplates"
       />
     </div>
+    <div class="include-test">
+      <FormCheckbox
+        label="Include a test file for the generated code"
+        saveKey="include_test"
+      />
+    </div>
   </div>
 </template>
 
 <script>
 import FormSelect from './FormSelect.vue'
+import FormCheckbox from './FormCheckbox.vue'
 import templates from '../templates/templates.json'
 import { store, fetchTemplates } from '../store.js'
 
 export default {
-  components: { FormSelect },
+  components: { FormSelect, FormCheckbox },
   setup() {
     const templateLabel = 'Choose A Template'
     const templateOptions = Object.keys(templates)
diff --git a/src/store.js b/src/store.js
index 106660b3..4804dcd2 100644
--- a/src/store.js
+++ b/src/store.js
@@ -40,7 +40,8 @@ export const msg = reactive({
 export const store = reactive({
   code: {},
   config: {
-    template: ''
+    template: '',
+    include_test: false
   }
 })
 
@@ -61,6 +62,10 @@ export function genCode() {
   const currentFiles = files[store.config.template]
   if (currentFiles && Object.keys(currentFiles).length) {
     for (const file in currentFiles) {
+      if (!store.config.include_test && file === 'test_all.py') {
+        delete store.code['test_all.py']
+        continue
+      }
       store.code[file] = ejs
         .render(currentFiles[file], store.config)
         .replaceAll(/(\n\n\n\n)+/gi, '\n')
diff --git a/src/templates/template-common/requirements.txt b/src/templates/template-common/requirements.txt
index 4bbe4290..52dc19aa 100644
--- a/src/templates/template-common/requirements.txt
+++ b/src/templates/template-common/requirements.txt
@@ -12,3 +12,7 @@ pyyaml
 #:::= it.logger :::#
 
 #::: } :::#
+
+#::: if (it.include_tests) { :::#
+pytest
+#::: } :::#
diff --git a/src/templates/template-vision-classification/requirements.txt b/src/templates/template-vision-classification/requirements.txt
index 4bbe4290..52dc19aa 100644
--- a/src/templates/template-vision-classification/requirements.txt
+++ b/src/templates/template-vision-classification/requirements.txt
@@ -12,3 +12,7 @@ pyyaml
 #:::= it.logger :::#
 
 #::: } :::#
+
+#::: if (it.include_tests) { :::#
+pytest
+#::: } :::#
diff --git a/src/templates/template-vision-classification/test_all.py b/src/templates/template-vision-classification/test_all.py
new file mode 100644
index 00000000..3313d3df
--- /dev/null
+++ b/src/templates/template-vision-classification/test_all.py
@@ -0,0 +1,54 @@
+import os
+from argparse import Namespace
+from typing import Iterable
+
+import ignite.distributed as idist
+import pytest
+import torch
+from data import setup_data
+from torch import Tensor, nn, optim
+from torch.utils.data.dataloader import DataLoader
+from trainers import setup_evaluator
+
+
+def set_up():
+    model = nn.Linear(1, 1)
+    optimizer = optim.Adam(model.parameters())
+    device = idist.device()
+    loss_fn = nn.MSELoss()
+    batch = [torch.tensor([1.0]), torch.tensor([1.0])]
+
+    return model, optimizer, device, loss_fn, batch
+
+
+@pytest.mark.skipif(
+    os.getenv("RUN_SLOW_TESTS", 0) == 0, reason="Skip slow tests"
+)
+def test_setup_data():
+    config = Namespace(
+        data_path="~/data", train_batch_size=1, eval_batch_size=1, num_workers=0
+    )
+    dataloader_train, dataloader_eval = setup_data(config)
+
+    assert isinstance(dataloader_train, DataLoader)
+    assert isinstance(dataloader_eval, DataLoader)
+    train_batch = next(iter(dataloader_train))
+    assert isinstance(train_batch, Iterable)
+    assert isinstance(train_batch[0], Tensor)
+    assert isinstance(train_batch[1], Tensor)
+    assert train_batch[0].ndim == 4
+    assert train_batch[1].ndim == 1
+    eval_batch = next(iter(dataloader_eval))
+    assert isinstance(eval_batch, Iterable)
+    assert isinstance(eval_batch[0], Tensor)
+    assert isinstance(eval_batch[1], Tensor)
+    assert eval_batch[0].ndim == 4
+    assert eval_batch[1].ndim == 1
+
+
+def test_setup_evaluator():
+    model, _, device, _, batch = set_up()
+    config = Namespace(use_amp=False)
+    evaluator = setup_evaluator(config, model, device)
+    evaluator.run([batch, batch])
+    assert isinstance(evaluator.state.output, tuple)
diff --git a/src/templates/template-vision-dcgan/requirements.txt b/src/templates/template-vision-dcgan/requirements.txt
index 4bbe4290..52dc19aa 100644
--- a/src/templates/template-vision-dcgan/requirements.txt
+++ b/src/templates/template-vision-dcgan/requirements.txt
@@ -12,3 +12,7 @@ pyyaml
 #:::= it.logger :::#
 
 #::: } :::#
+
+#::: if (it.include_tests) { :::#
+pytest
+#::: } :::#
diff --git a/src/templates/template-vision-dcgan/test_all.py b/src/templates/template-vision-dcgan/test_all.py
new file mode 100644
index 00000000..8214f397
--- /dev/null
+++ b/src/templates/template-vision-dcgan/test_all.py
@@ -0,0 +1,71 @@
+import os
+from argparse import Namespace
+from numbers import Number
+from typing import Iterable
+
+import ignite.distributed as idist
+import pytest
+import torch
+from data import setup_data
+from model import Discriminator, Generator
+from torch import Tensor, nn, optim
+from torch.utils.data.dataloader import DataLoader
+from trainers import setup_trainer
+
+
+def set_up():
+    model = nn.Linear(1, 1)
+    optimizer = optim.Adam(model.parameters())
+    device = idist.device()
+    loss_fn = nn.MSELoss()
+    batch = [torch.tensor([1.0]), torch.tensor([1.0])]
+
+    return model, optimizer, device, loss_fn, batch
+
+
+@pytest.mark.skipif(
+    os.getenv("RUN_SLOW_TESTS", 0) == 0, reason="Skip slow tests"
+)
+def test_setup_data():
+    config = Namespace(
+        data_path="~/data", train_batch_size=1, eval_batch_size=1, num_workers=0
+    )
+    dataloader_train, dataloader_eval, _ = setup_data(config)
+
+    assert isinstance(dataloader_train, DataLoader)
+    assert isinstance(dataloader_eval, DataLoader)
+    train_batch = next(iter(dataloader_train))
+    assert isinstance(train_batch, Iterable)
+    assert isinstance(train_batch[0], Tensor)
+    assert isinstance(train_batch[1], Tensor)
+    assert train_batch[0].ndim == 4
+    assert train_batch[1].ndim == 1
+    eval_batch = next(iter(dataloader_eval))
+    assert isinstance(eval_batch, Iterable)
+    assert isinstance(eval_batch[0], Tensor)
+    assert isinstance(eval_batch[1], Tensor)
+    assert eval_batch[0].ndim == 4
+    assert eval_batch[1].ndim == 1
+
+
+def test_models():
+    model_G = Generator(100, 64, 3)
+    model_D = Discriminator(3, 64)
+    x = torch.rand([1, 100, 32, 32])
+    y = model_G(x)
+    y.sum().backward()
+    z = model_D(y)
+    assert y.shape == torch.Size([1, 3, 560, 560])
+    assert z.shape == torch.Size([1024])
+    assert isinstance(model_D, nn.Module)
+    assert isinstance(model_G, nn.Module)
+
+
+def test_setup_trainer():
+    model, optimizer, device, loss_fn, batch = set_up()
+    config = Namespace(use_amp=False, train_batch_size=2, z_dim=100)
+    trainer = setup_trainer(
+        config, model, model, optimizer, optimizer, loss_fn, device
+    )
+    trainer.run([batch, batch])
+    assert isinstance(trainer.state.output, dict)
diff --git a/src/templates/template-vision-segmentation/data.py b/src/templates/template-vision-segmentation/data.py
index acd16add..79c403de 100644
--- a/src/templates/template-vision-segmentation/data.py
+++ b/src/templates/template-vision-segmentation/data.py
@@ -133,7 +133,7 @@ def setup_data(config: Namespace):
     dataloader_eval = idist.auto_dataloader(
         dataset_eval,
         shuffle=False,
-        batch_size=config.train_batch_size,
+        batch_size=config.eval_batch_size,
         num_workers=config.num_workers,
         drop_last=False,
     )
diff --git a/src/templates/template-vision-segmentation/requirements.txt b/src/templates/template-vision-segmentation/requirements.txt
index 22d3c3d5..054bac44 100644
--- a/src/templates/template-vision-segmentation/requirements.txt
+++ b/src/templates/template-vision-segmentation/requirements.txt
@@ -14,3 +14,7 @@ image_dataset_viz
 #:::= it.logger :::#
 
 #::: } :::#
+
+#::: if (it.include_tests) { :::#
+pytest
+#::: } :::#
diff --git a/src/templates/template-vision-segmentation/test_all.py b/src/templates/template-vision-segmentation/test_all.py
new file mode 100644
index 00000000..6859c879
--- /dev/null
+++ b/src/templates/template-vision-segmentation/test_all.py
@@ -0,0 +1,32 @@
+import os
+from argparse import Namespace
+
+import pytest
+from data import setup_data
+from torch import Tensor
+from torch.utils.data.dataloader import DataLoader
+
+
+@pytest.mark.skipif(
+    os.getenv("RUN_SLOW_TESTS", 0) == 0, reason="Skip slow tests"
+)
+def test_setup_data():
+    config = Namespace(
+        data_path="~/data", train_batch_size=1, eval_batch_size=1, num_workers=0
+    )
+    dataloader_train, dataloader_eval = setup_data(config)
+
+    assert isinstance(dataloader_train, DataLoader)
+    assert isinstance(dataloader_eval, DataLoader)
+    train_batch = next(iter(dataloader_train))
+    assert isinstance(train_batch, dict)
+    assert isinstance(train_batch["image"], Tensor)
+    assert isinstance(train_batch["mask"], Tensor)
+    assert train_batch["image"].ndim == 4
+    assert train_batch["mask"].ndim == 3
+    eval_batch = next(iter(dataloader_eval))
+    assert isinstance(eval_batch, dict)
+    assert isinstance(eval_batch["image"], Tensor)
+    assert isinstance(eval_batch["mask"], Tensor)
+    assert eval_batch["image"].ndim == 4
+    assert eval_batch["mask"].ndim == 3
diff --git a/src/templates/templates.json b/src/templates/templates.json
index d1d57867..248b029d 100644
--- a/src/templates/templates.json
+++ b/src/templates/templates.json
@@ -7,6 +7,7 @@
     "model.py",
     "trainers.py",
     "utils.py",
+    "test_all.py",
     "requirements.txt"
   ],
   "template-vision-dcgan": [
@@ -17,6 +18,7 @@
     "model.py",
     "trainers.py",
     "utils.py",
+    "test_all.py",
     "requirements.txt"
   ],
   "template-vision-segmentation": [
@@ -27,6 +29,7 @@
     "trainers.py",
     "utils.py",
     "vis.py",
+    "test_all.py",
     "requirements.txt"
   ]
 }