diff --git a/.circleci/config.yml b/.circleci/config.yml
index 22d6981b..8b014494 100755
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -10,10 +10,10 @@ references:
     run:
       name: Install Dependences
       command: |
-        pip install "$TORCH_VERSION" --user
-        pip install -r requirements.txt --user
-        sudo pip install pytest pytest-cov pytest-flake8
-        pip install -r ./tests/requirements.txt --user
+        pip install "$TORCH_VERSION"
+        pip install -r requirements.txt -q
+        sudo pip install pytest pytest-cov pytest-flake8 -q
+        pip install -r ./tests/requirements.txt -q
 
   tests: &tests
    run:
@@ -21,7 +21,7 @@ references:
      command: |
        python --version ; pip --version ; pip list
        py.test pytorch_lightning tests -v --doctest-modules --junitxml=test-reports/pytest_junit.xml
-     no_output_timeout: 15m
+     no_output_timeout: 30m
 
   examples: &examples
    run:
@@ -30,7 +30,7 @@ references:
        pip install -r ./pl_examples/requirements.txt --user
        python --version ; pip --version ; pip list
        py.test pl_examples -v --doctest-modules --junitxml=test-reports/pytest_junit.xml
-     no_output_timeout: 15m
+     no_output_timeout: 20m
 
   install_pkg: &install_pkg
     run:
@@ -53,8 +53,10 @@ references:
    run:
      name: Formatting
      command: |
-       python --version ; pip --version ; pip list
-       flake8
+       python --version ; pip --version
+       sudo pip install flake8 -q
+       pip list
+       flake8 .
 
   make_docs: &make_docs
    run:
@@ -82,45 +84,56 @@ jobs:
         - TORCH_VERSION: "torch"
     steps:
       - checkout
-      - *install_deps
       - *format
 
   PyTorch:
     docker:
-      - image: circleci/python:3.7
+      - image: circleci/python:3.6
     environment:
       - TORCH_VERSION: "torch"
     steps: &steps
       - checkout
+      #- restore_cache:
+      #    keys:
+      #      # when lock file changes, use increasingly general patterns to restore cache
+      #      - pip-packages--{{ .Environment.CIRCLE_JOB }}
+      #      - pip-packages--
       - *install_deps
+      #- save_cache:
+      #    key: pip-packages--{{ .Environment.CIRCLE_JOB }}
+      #    paths:
+      #      # this path depends on where pipenv creates a virtualenv
+      #      - "~/.cache/pip"
+      #      - "/usr/local/lib/python3.6/site-packages"
+      #      - "/usr/local/lib/site-python"
       - *tests
       - store_test_results:
           path: test-reports
       - store_artifacts:
           path: test-reports
 
-  PyTorch-v1.1:
+  PyTorch-v1_1:
     docker:
       - image: circleci/python:3.6
     environment:
       - TORCH_VERSION: "torch>=1.1, <1.2"
     steps: *steps
 
-  PyTorch-v1.2:
+  PyTorch-v1_2:
     docker:
       - image: circleci/python:3.6
     environment:
       - TORCH_VERSION: "torch>=1.2, <1.3"
     steps: *steps
 
-  PyTorch-v1.3:
+  PyTorch-v1_3:
     docker:
       - image: circleci/python:3.6
     environment:
       - TORCH_VERSION: "torch>=1.3, <1.4"
     steps: *steps
 
-  PyTorch-v1.4:
+  PyTorch-v1_4:
     docker:
       - image: circleci/python:3.6
     environment:
@@ -145,6 +158,8 @@ jobs:
       - *create_pkg
       - *install_pkg
 
+#orbs:
+#  python: circleci/python@0.2.1
 
 workflows:
   version: 2
@@ -152,9 +167,9 @@ workflows:
     jobs:
       - Formatting
       - Build-Docs
-      - PyTorch-v1.1
-      - PyTorch-v1.2
-      - PyTorch-v1.3
-      - PyTorch-v1.4
+      - PyTorch-v1_1
+      - PyTorch-v1_2
+      - PyTorch-v1_3
+      - PyTorch-v1_4
       - Install-pkg
       - Examples
diff --git a/.drone.yml b/.drone.yml
index d1dbc4c1..90c8873b 100644
--- a/.drone.yml
+++ b/.drone.yml
@@ -7,22 +7,42 @@ name: torch-GPU
 steps:
 - name: testing
   image: pytorch/pytorch:1.4-cuda10.1-cudnn7-runtime
+
   environment:
     SLURM_LOCALID: 0
     CODECOV_TOKEN:
       from_secret: codecov_token
+
+  #volumes:
+  #  # Mount pip cache from host
+  #  - name: pip_cache
+  #    path: /opt/conda/lib/python3.7/site-packages
+
   commands:
     - python --version
     - pip install pip -U
     - pip --version
     - nvidia-smi
     - bash ./tests/install_AMP.sh
-    - pip install -r requirements.txt --user
-    - pip install coverage pytest pytest-cov pytest-flake8 codecov
-    - pip install -r ./tests/requirements.txt --user
+    - pip install -r requirements.txt --user -q
+    - pip install coverage pytest pytest-cov pytest-flake8 codecov -q
+    - pip install -r ./tests/requirements.txt --user -q
     - pip list
     - python -c "import torch ; print(' & '.join([torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())]) if torch.cuda.is_available() else 'only CPU')"
     - coverage run --source pytorch_lightning -m py.test pytorch_lightning tests benchmarks -v --doctest-modules #  --flake8
     - coverage report
     - codecov --token $CODECOV_TOKEN  # --pr $DRONE_PULL_REQUEST --build $DRONE_BUILD_NUMBER --branch $DRONE_BRANCH --commit $DRONE_COMMIT --tag $DRONE_TAG
     - python tests/collect_env_details.py
+
+trigger:
+  branch:
+    - master
+  event:
+    include:
+    - push
+    - pull_request
+
+#volumes:
+#  - name: pip_cache
+#    host:
+#      path: /tmp/cache/drone/pip
diff --git a/.github/workflows/ci-testing.yml b/.github/workflows/ci-testing.yml
index ab6d1317..6491ef02 100644
--- a/.github/workflows/ci-testing.yml
+++ b/.github/workflows/ci-testing.yml
@@ -28,7 +28,7 @@ jobs:
             requires: 'minimal'
 
     # Timeout: https://stackoverflow.com/a/59076067/4521646
-    timeout-minutes: 20
+    timeout-minutes: 30
     steps:
     - uses: actions/checkout@v2
     - name: Set up Python ${{ matrix.python-version }}
@@ -61,15 +61,15 @@ jobs:
     - name: Cache datasets
       uses: actions/cache@v1
       with:
-        path: tests/datasets # This path is specific to Ubuntu
+        path: tests/Datasets # This path is specific to Ubuntu
         # Look to see if there is a cache hit for the corresponding requirements file
         key: mnist-dataset
 
     - name: Install dependencies
       run: |
         # python -m pip install --upgrade --user pip
-        pip install -r requirements.txt -U -f https://download.pytorch.org/whl/torch_stable.html
-        pip install -r ./tests/requirements.txt
+        pip install -r requirements.txt -U -f https://download.pytorch.org/whl/torch_stable.html -q
+        pip install -r ./tests/requirements.txt -q
         # pip install tox coverage
         python --version
         pip --version
diff --git a/benchmarks/test_trainer_parity.py b/benchmarks/test_trainer_parity.py
index 37edf8f1..3cb601b3 100644
--- a/benchmarks/test_trainer_parity.py
+++ b/benchmarks/test_trainer_parity.py
@@ -8,9 +8,9 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torch.utils.data import DataLoader
 from torchvision import transforms
-from torchvision.datasets import MNIST
 
 from pytorch_lightning import Trainer, LightningModule
+from tests.base.datasets import TestingMNIST
 
 
 class ParityMNIST(LightningModule):
@@ -41,8 +41,11 @@ class ParityMNIST(LightningModule):
         return torch.optim.Adam(self.parameters(), lr=0.02)
 
     def train_dataloader(self):
-        return DataLoader(MNIST(os.getcwd(), train=True, download=True,
-                                transform=transforms.ToTensor()), batch_size=32)
+        return DataLoader(TestingMNIST(train=True,
+                                       download=True,
+                                       num_samples=500,
+                                       digits=list(range(5))),
+                          batch_size=128)
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU machine")
diff --git a/tests/__init__.py b/tests/__init__.py
index e69de29b..32e384ac 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -0,0 +1,3 @@
+import os
+
+TEST_ROOT = os.path.dirname(__file__)
diff --git a/tests/base/datasets.py b/tests/base/datasets.py
index 0fbdfae6..8231c093 100644
--- a/tests/base/datasets.py
+++ b/tests/base/datasets.py
@@ -1,12 +1,17 @@
 import logging
 import os
 import urllib.request
-from typing import Tuple
+from typing import Tuple, Optional, Sequence
 
 import torch
 from torch import Tensor
 from torch.utils.data import Dataset
 
+from tests import TEST_ROOT
+
+#: local path to test datasets
+PATH_DATASETS = os.path.join(TEST_ROOT, 'Datasets')
+
 
 class MNIST(Dataset):
     """
@@ -25,6 +30,13 @@ class MNIST(Dataset):
         download: If true, downloads the dataset from the internet and
             puts it in root directory. If dataset is already downloaded, it is not
             downloaded again.
+
+    Examples:
+        >>> dataset = MNIST(download=True)
+        >>> len(dataset)
+        60000
+        >>> torch.bincount(dataset.targets)
+        tensor([5923, 6742, 5958, 6131, 5842, 5421, 5918, 6265, 5851, 5949])
     """
 
     RESOURCES = (
@@ -34,21 +46,22 @@ class MNIST(Dataset):
 
     TRAIN_FILE_NAME = 'training.pt'
     TEST_FILE_NAME = 'test.pt'
+    cache_folder_name = 'complete'
 
-    def __init__(self, root: str, train: bool = True, normalize: tuple = (0.5, 1.0), download: bool = False):
-        super(MNIST, self).__init__()
+    def __init__(self, root: str = PATH_DATASETS, train: bool = True,
+                 normalize: tuple = (0.5, 1.0), download: bool = False):
+        super().__init__()
         self.root = root
         self.train = train  # training set or test set
         self.normalize = normalize
 
-        if download:
-            self.download()
+        self.prepare_data(download)
 
-        if not self._check_exists():
+        if not self._check_exists(self.cached_folder_path):
             raise RuntimeError('Dataset not found.')
 
         data_file = self.TRAIN_FILE_NAME if self.train else self.TEST_FILE_NAME
-        self.data, self.targets = torch.load(os.path.join(self.processed_folder, data_file))
+        self.data, self.targets = torch.load(os.path.join(self.cached_folder_path, data_file))
 
     def __getitem__(self, idx: int) -> Tuple[Tensor, int]:
         img = self.data[idx].float().unsqueeze(0)
@@ -63,25 +76,30 @@ class MNIST(Dataset):
         return len(self.data)
 
     @property
-    def processed_folder(self) -> str:
-        return os.path.join(self.root, 'MNIST', 'processed')
+    def cached_folder_path(self) -> str:
+        return os.path.join(self.root, 'MNIST', self.cache_folder_name)
 
-    def _check_exists(self) -> bool:
-        train_file = os.path.join(self.processed_folder, self.TRAIN_FILE_NAME)
-        test_file = os.path.join(self.processed_folder, self.TEST_FILE_NAME)
-        return os.path.isfile(train_file) and os.path.isfile(test_file)
+    def _check_exists(self, data_folder: str) -> bool:
+        existing = True
+        for fname in (self.TRAIN_FILE_NAME, self.TEST_FILE_NAME):
+            existing = existing and os.path.isfile(os.path.join(data_folder, fname))
+        return existing
 
-    def download(self) -> None:
-        """Download the MNIST data if it doesn't exist in processed_folder already."""
+    def prepare_data(self, download: bool):
+        if download:
+            self._download(self.cached_folder_path)
 
-        if self._check_exists():
+    def _download(self, data_folder: str) -> None:
+        """Download the MNIST data if it doesn't exist in cached_folder_path already."""
+
+        if self._check_exists(data_folder):
             return
 
-        os.makedirs(self.processed_folder, exist_ok=True)
+        os.makedirs(data_folder, exist_ok=True)
 
         for url in self.RESOURCES:
             logging.info(f'Downloading {url}')
-            fpath = os.path.join(self.processed_folder, os.path.basename(url))
+            fpath = os.path.join(data_folder, os.path.basename(url))
             urllib.request.urlretrieve(url, fpath)
 
 
@@ -94,14 +112,73 @@ def normalize_tensor(tensor: Tensor, mean: float = 0.0, std: float = 1.0) -> Ten
 
 
 class TestingMNIST(MNIST):
+    """Constrain image dataset
+
+    Args:
+        root: Root directory of dataset where ``MNIST/processed/training.pt``
+            and  ``MNIST/processed/test.pt`` exist.
+        train: If ``True``, creates dataset from ``training.pt``,
+            otherwise from ``test.pt``.
+        normalize: mean and std deviation of the MNIST dataset.
+        download: If true, downloads the dataset from the internet and
+            puts it in root directory. If dataset is already downloaded, it is not
+            downloaded again.
+        num_samples: number of examples per selected class/digit
+        digits: list selected MNIST digits/classes
+
+    Examples:
+        >>> dataset = TestingMNIST(download=True)
+        >>> len(dataset)
+        300
+        >>> sorted(set([d.item() for d in dataset.targets]))
+        [0, 1, 2]
+        >>> torch.bincount(dataset.targets)
+        tensor([100, 100, 100])
+    """
+
+    def __init__(self, root: str = PATH_DATASETS, train: bool = True,
+                 normalize: tuple = (0.5, 1.0), download: bool = False,
+                 num_samples: int = 100, digits: Optional[Sequence] = (0, 1, 2)):
+
+        # number of examples per class
+        self.num_samples = num_samples
+        # take just a subset of MNIST dataset
+        self.digits = digits if digits else list(range(10))
+
+        self.cache_folder_name = 'digits-' + '-'.join(str(d) for d in sorted(self.digits)) \
+                                 + f'_nb-{self.num_samples}'
 
-    def __init__(self, root, train=True, normalize=(0.5, 1.0), download=False, num_samples=8000):
         super().__init__(
             root,
             train=train,
             normalize=normalize,
             download=download
         )
-        # take just a subset of MNIST dataset
-        self.data = self.data[:num_samples]
-        self.targets = self.targets[:num_samples]
+
+    @staticmethod
+    def _prepare_subset(full_data: torch.Tensor, full_targets: torch.Tensor,
+                        num_samples: int, digits: Sequence):
+        classes = {d: 0 for d in digits}
+        indexes = []
+        for idx, target in enumerate(full_targets):
+            label = target.item()
+            if classes.get(label, float('inf')) >= num_samples:
+                continue
+            indexes.append(idx)
+            classes[label] += 1
+            if all(classes[k] >= num_samples for k in classes):
+                break
+        data = full_data[indexes]
+        targets = full_targets[indexes]
+        return data, targets
+
+    def prepare_data(self, download: bool) -> None:
+        if self._check_exists(self.cached_folder_path):
+            return
+        if download:
+            self._download(super().cached_folder_path)
+
+        for fname in (self.TRAIN_FILE_NAME, self.TEST_FILE_NAME):
+            data, targets = torch.load(os.path.join(super().cached_folder_path, fname))
+            data, targets = self._prepare_subset(data, targets, self.num_samples, self.digits)
+            torch.save((data, targets), os.path.join(self.cached_folder_path, fname))
diff --git a/tests/base/debug.py b/tests/base/debug.py
index b601982d..33be2899 100644
--- a/tests/base/debug.py
+++ b/tests/base/debug.py
@@ -3,7 +3,7 @@ from torch.nn import functional as F
 from torch.utils.data import DataLoader
 
 import pytorch_lightning as pl
-from tests.base.datasets import MNIST
+from tests.base.datasets import TestingMNIST
 
 
 # from test_models import assert_ok_test_acc, load_model, \
@@ -42,10 +42,10 @@ class CoolModel(pl.LightningModule):
         return [torch.optim.Adam(self.parameters(), lr=0.02)]
 
     def train_dataloader(self):
-        return DataLoader(MNIST('path/to/save', train=True), batch_size=32)
+        return DataLoader(TestingMNIST(train=True, num_samples=100), batch_size=16)
 
     def val_dataloader(self):
-        return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
+        return DataLoader(TestingMNIST(train=False, num_samples=50), batch_size=16)
 
     def test_dataloader(self):
-        return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
+        return DataLoader(TestingMNIST(train=False, num_samples=50), batch_size=16)
diff --git a/tests/base/models.py b/tests/base/models.py
index b796b60c..50419e52 100644
--- a/tests/base/models.py
+++ b/tests/base/models.py
@@ -38,7 +38,7 @@ class DictHparamsModel(LightningModule):
         return torch.optim.Adam(self.parameters(), lr=0.02)
 
     def train_dataloader(self):
-        return DataLoader(TestingMNIST(os.getcwd(), train=True, download=True), batch_size=32)
+        return DataLoader(TestingMNIST(train=True, download=True), batch_size=16)
 
 
 class TestModelBase(LightningModule):
@@ -137,13 +137,11 @@ class TestModelBase(LightningModule):
         return [optimizer], [scheduler]
 
     def prepare_data(self):
-        _ = TestingMNIST(root=self.hparams.data_root, train=True,
-                         download=True, num_samples=2000)
+        _ = TestingMNIST(root=self.hparams.data_root, train=True, download=True)
 
     def _dataloader(self, train):
         # init data generators
-        dataset = TestingMNIST(root=self.hparams.data_root, train=train,
-                               download=False, num_samples=2000)
+        dataset = TestingMNIST(root=self.hparams.data_root, train=train, download=False)
 
         # when using multi-node we need to add the datasampler
         batch_size = self.hparams.batch_size
diff --git a/tests/base/utils.py b/tests/base/utils.py
index 9eb42a5c..e5639362 100644
--- a/tests/base/utils.py
+++ b/tests/base/utils.py
@@ -10,6 +10,7 @@ from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import ModelCheckpoint
 from pytorch_lightning.loggers import TestTubeLogger, TensorBoardLogger
 from tests.base import LightningTestModel
+from tests.base.datasets import PATH_DATASETS
 
 # generate a list of random seeds for each test
 RANDOM_PORTS = list(np.random.randint(12000, 19000, 1000))
@@ -99,7 +100,7 @@ def get_default_hparams(continue_training=False, hpc_exp_number=0):
         'in_features': 28 * 28,
         'learning_rate': 0.001 * 8,
         'optimizer_name': 'adam',
-        'data_root': os.path.join(tests_dir, 'datasets'),
+        'data_root': PATH_DATASETS,
         'out_features': 10,
         'hidden_dim': 1000,
     }
@@ -180,7 +181,7 @@ def load_model_from_checkpoint(root_weights_dir, module_class=LightningTestModel
     return trained_model
 
 
-def run_prediction(dataloader, trained_model, dp=False, min_acc=0.35):
+def run_prediction(dataloader, trained_model, dp=False, min_acc=0.5):
     # run prediction on 1 batch
     for batch in dataloader:
         break
@@ -205,7 +206,7 @@ def run_prediction(dataloader, trained_model, dp=False, min_acc=0.35):
     assert acc >= min_acc, f"This model is expected to get > {min_acc} in test set (it got {acc})"
 
 
-def assert_ok_model_acc(trainer, key='test_acc', thr=0.4):
+def assert_ok_model_acc(trainer, key='test_acc', thr=0.5):
     # this model should get 0.80+ acc
     acc = trainer.training_tqdm_dict[key]
     assert acc > thr, f"Model failed to get expected {thr} accuracy. {key} = {acc}"
diff --git a/tests/datasets/__init__.py b/tests/datasets/__init__.py
deleted file mode 100644
index e69de29b..00000000
diff --git a/tests/models/test_cpu.py b/tests/models/test_cpu.py
index 60baec3f..5ce2d195 100644
--- a/tests/models/test_cpu.py
+++ b/tests/models/test_cpu.py
@@ -55,7 +55,7 @@ def test_lbfgs_cpu_model(tmpdir):
     )
 
     model, hparams = tutils.get_default_model(lbfgs=True)
-    tutils.run_model_test_no_loggers(trainer_options, model, min_acc=0.30)
+    tutils.run_model_test_no_loggers(trainer_options, model, min_acc=0.5)
 
 
 def test_default_logger_callbacks_cpu_model(tmpdir):
@@ -113,7 +113,7 @@ def test_running_test_after_fitting(tmpdir):
     trainer.test()
 
     # test we have good test accuracy
-    tutils.assert_ok_model_acc(trainer, thr=0.35)
+    tutils.assert_ok_model_acc(trainer, thr=0.5)
 
 
 def test_running_test_without_val(tmpdir):
diff --git a/tests/models/test_restore.py b/tests/models/test_restore.py
index d1f5e6ec..37a04ec5 100644
--- a/tests/models/test_restore.py
+++ b/tests/models/test_restore.py
@@ -166,7 +166,7 @@ def test_running_test_pretrained_model_dp(tmpdir):
 
     trainer_options = dict(
         show_progress_bar=True,
-        max_epochs=4,
+        max_epochs=2,
         train_percent_check=0.4,
         val_percent_check=0.2,
         checkpoint_callback=checkpoint,
@@ -203,7 +203,7 @@ def test_dp_resume(tmpdir):
 
     trainer_options = dict(
         show_progress_bar=True,
-        max_epochs=3,
+        max_epochs=1,
         gpus=2,
         distributed_backend='dp',
     )
diff --git a/tests/trainer/test_callbacks.py b/tests/trainer/test_callbacks.py
index 5d0c3961..ba398320 100644
--- a/tests/trainer/test_callbacks.py
+++ b/tests/trainer/test_callbacks.py
@@ -173,11 +173,11 @@ def test_early_stopping_without_val_step(tmpdir):
         default_save_path=tmpdir,
         early_stop_callback=stopping,
         overfit_pct=0.20,
-        max_epochs=10,
+        max_epochs=5,
     )
 
     trainer = Trainer(**trainer_options)
     result = trainer.fit(model)
 
     assert result == 1, 'training failed to complete'
-    assert trainer.current_epoch < trainer.max_epochs - 1
+    assert trainer.current_epoch < trainer.max_epochs
diff --git a/tests/trainer/test_optimizers.py b/tests/trainer/test_optimizers.py
index 213529a9..b07cbd4f 100644
--- a/tests/trainer/test_optimizers.py
+++ b/tests/trainer/test_optimizers.py
@@ -140,7 +140,7 @@ def test_multi_optimizer_with_scheduling_stepping(tmpdir):
     adjusted_lr2 = adjusted_lr2[0]
 
     # Called ones after end of epoch
-    assert init_lr * (0.1)**3 == adjusted_lr1, \
+    assert init_lr * 0.1 ** 0 == adjusted_lr1, \
         'lr for optimizer 1 not adjusted correctly'
     # Called every 3 steps, meaning for 1 epoch of 11 batches, it is called 3 times
     assert init_lr * 0.1 == adjusted_lr2, \
diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
index 12ea5941..923a4f83 100644
--- a/tests/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -34,7 +34,7 @@ def test_hparams_save_load(tmpdir):
     # logger file to get meta
     trainer_options = dict(
         default_save_path=tmpdir,
-        max_epochs=2,
+        max_epochs=1,
     )
 
     # fit model
@@ -197,7 +197,7 @@ def test_gradient_accumulation_scheduling(tmpdir):
     trainer = Trainer(accumulate_grad_batches=schedule,
                       train_percent_check=0.1,
                       val_percent_check=0.1,
-                      max_epochs=4,
+                      max_epochs=2,
                       default_save_path=tmpdir)
 
     # for the test
@@ -360,10 +360,10 @@ def test_resume_from_checkpoint_epoch_restored(tmpdir):
         state = torch.load(check)
 
         # Resume training
-        trainer_options['max_epochs'] = 4
+        trainer_options['max_epochs'] = 2
         new_trainer = Trainer(**trainer_options, resume_from_checkpoint=check)
         new_trainer.fit(next_model)
-        assert state['global_step'] + next_model.num_batches_seen == training_batches * 4
+        assert state['global_step'] + next_model.num_batches_seen == training_batches * trainer_options['max_epochs']
 
 
 def _init_steps_model():
@@ -372,7 +372,7 @@ def _init_steps_model():
     model, _ = tutils.get_default_model()
 
     # define train epoch to 5% of data
-    train_percent = 0.05
+    train_percent = 0.5
     # get number of samples in 1 epoch
     num_train_samples = math.floor(len(model.train_dataloader()) * train_percent)
 
@@ -389,7 +389,7 @@ def test_trainer_max_steps_and_epochs(tmpdir):
     # define less train steps than epochs
     trainer_options.update(dict(
         default_save_path=tmpdir,
-        max_epochs=5,
+        max_epochs=3,
         max_steps=num_train_samples + 10
     ))
 
@@ -413,8 +413,8 @@ def test_trainer_max_steps_and_epochs(tmpdir):
     assert result == 1, "Training did not complete"
 
     # check training stopped at max_epochs
-    assert trainer.global_step == num_train_samples * trainer.max_epochs \
-        and trainer.current_epoch == trainer.max_epochs - 1, "Model did not stop at max_epochs"
+    assert trainer.global_step == num_train_samples * trainer.max_epochs
+    assert trainer.current_epoch == trainer.max_epochs - 1, "Model did not stop at max_epochs"
 
 
 def test_trainer_min_steps_and_epochs(tmpdir):
@@ -427,7 +427,7 @@ def test_trainer_min_steps_and_epochs(tmpdir):
         early_stop_callback=EarlyStopping(monitor='val_loss', min_delta=1.0),
         val_check_interval=2,
         min_epochs=1,
-        max_epochs=10
+        max_epochs=5
     ))
 
     # define less min steps than 1 epoch