diff --git a/fastai_contrib/data.py b/fastai_contrib/data.py
index e14706d..274d1f6 100644
--- a/fastai_contrib/data.py
+++ b/fastai_contrib/data.py
@@ -37,9 +37,17 @@ class LanguageModelLoader(): # copy of the original LanguageModelLoader
             itr += 1
             yield res
 
-    def __len__(self) -> int: return (self.n-1) // self.bptt
+    def __len__(self) -> int: return int(math.ceil((self.n-1) / self.bptt)) # so that it is always at least 1
     def __getattr__(self,k:str)->Any: return getattr(self.dataset, k)
 
+    @property
+    def batch_size(self):
+        return self.bs
+
+    @batch_size.setter
+    def batch_size(self, v):
+        self.bs = v
+
     def batchify(self, data:np.ndarray) -> LongTensor:
         "Split the corpus `data` in batches."
         nb = data.shape[0] // self.bs
diff --git a/tests/test_end_to_end.py b/tests/test_end_to_end.py
index 355ffbf..9b80a6b 100644
--- a/tests/test_end_to_end.py
+++ b/tests/test_end_to_end.py
@@ -11,10 +11,8 @@ from fastai_contrib.utils import *
 It is a mixture of a pytest unit test and woven together to compose an end to end functional test. 
 """
 
-
-def delete_test_models():
-    data = get_data_folder()
-
+import fastai.core
+fastai.core.turn_off_parallel_execution=True
 
 def copy_head(src_fn, dst_fn, n=1000):
     with src_fn.open("r") as s, dst_fn.open("w") as d:
@@ -35,7 +33,7 @@ def get_test_data():
     test_wt.mkdir(exist_ok=True, parents=True)
     test_imdb.mkdir(exist_ok=True, parents=True)
 
-    sz=10
+    sz=1
     # we use the same text to see if models can overfit
     copy_head(wt / 'en.wiki.train.tokens', test_wt / 'en.wiki.train.tokens', n=10*sz)
     copy_head(wt / 'en.wiki.train.tokens', test_wt / 'en.wiki.valid.tokens', n=6*sz)
@@ -71,11 +69,12 @@ def test_ulmfit_default_end_to_end():
         cuda_id=cuda_id,
         fine_tune=True,
         max_vocab=1000,
-        bs=2, bptt=70, name=lm_name + '-imdb-clas',
+        num_lm_epochs=0,
+        bs=4, # minimum size is 4 otherwise it somewhere becomes 1 and fit stops working
+        bptt=70,
+        name=lm_name + '-imdb-clas',
         dataset='imdb')
 
-    delete_test_models()
-
 
 def test_ulmfit_sentencepiece_end_to_end():
     """ Test ulmfit with sentencepiece tokenizer on small wikipedia dataset.
@@ -89,8 +88,8 @@ def test_ulmfit_sentencepiece_end_to_end():
         cuda_id=cuda_id,
         qrnn=True,
         subword=True,
-        max_vocab=1000,
-        bs=80,
+        max_vocab=100,
+        bs=2,
         num_epochs=1,
         name=lm_name,
     )
@@ -100,8 +99,7 @@ def test_ulmfit_sentencepiece_end_to_end():
     # NOTE: ds_pct is not available for sentencepiece -- tests are on the complete dataset
     #       sentencepiece for finetuning/classification is currently not implemented
 
-    delete_test_models()
-
 
 if __name__ == "__main__":
     fire.Fire()  # allows using all functions via CLI e.g. python utils.py prepare_imdb aclImdb.tgz
+
diff --git a/ulmfit/pretrain_lm.py b/ulmfit/pretrain_lm.py
index 5b7098a..aaca9e7 100644
--- a/ulmfit/pretrain_lm.py
+++ b/ulmfit/pretrain_lm.py
@@ -74,7 +74,9 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, subword=False, max_vo
 
         sp = get_sentencepiece(dir_path, trn_path, name, vocab_size=max_vocab)
 
-        data_lm = TextLMDataBunch.from_csv(dir_path, 'train.csv', **sp)
+        lm_type = contrib_data.LanguageModelType.BiLM if bidir else  contrib_data.LanguageModelType.FwdLM
+
+        data_lm = TextLMDataBunch.from_csv(dir_path, 'train.csv', **sp, bs=bs, bptt=bptt, lm_type=lm_type)
         itos = data_lm.train_ds.vocab.itos
         stoi = data_lm.train_ds.vocab.stoi
     else:
diff --git a/ulmfit/train_clas.py b/ulmfit/train_clas.py
index 7b617dd..7b88963 100644
--- a/ulmfit/train_clas.py
+++ b/ulmfit/train_clas.py
@@ -19,7 +19,7 @@ from pathlib import Path
 
 
 def new_train_clas(data_dir, lang='en', cuda_id=0, pretrain_name='wt103', model_dir='models',
-                   qrnn=False,
+                   qrnn=False, num_lm_epochs=10,
                    fine_tune=True, max_vocab=60000, bs=20, bptt=70, name='imdb-clas',
                    dataset='imdb', bidir=False, ds_pct=1.0, train=True):
     """
@@ -101,7 +101,7 @@ def new_train_clas(data_dir, lang='en', cuda_id=0, pretrain_name='wt103', model_
 
         learn.fit_one_cycle(1, 1e-2, moms=(0.8, 0.7))
         learn.unfreeze()
-        learn.fit_one_cycle(10, 1e-3, moms=(0.8, 0.7))
+        if num_lm_epochs > 0: learn.fit_one_cycle(num_lm_epochs, 1e-3, moms=(0.8, 0.7))
 
         # save encoder
         learn.save_encoder(lm_enc_finetuned)