diff --git a/fastai_contrib/learner.py b/fastai_contrib/learner.py
index bdcea7c..ea2588f 100644
--- a/fastai_contrib/learner.py
+++ b/fastai_contrib/learner.py
@@ -30,18 +30,22 @@ def bilm_learner(data:DataBunch, bptt:int=70, emb_sz:int=400, nh:int=1150, nl:in
 def bilm_text_classifier_learner(data: DataBunch, bptt: int = 70, max_len: int = 70 * 20, emb_sz: int = 400,
                             nh: int = 1150, nl: int = 3,
                             lin_ftrs: Collection[int] = None, ps: Collection[float] = None, pad_token: int = 1,
-                            drop_mult: float = 1., qrnn: bool = False, **kwargs) -> 'TextClassifierLearner':
+                            drop_mult: float = 1., qrnn: bool = False, bicls_head:str='BiPoolingLinearClassifier',  **kwargs) -> 'TextClassifierLearner':
     "Create a RNN classifier."
     dps = default_dropout['classifier'] * drop_mult
     if lin_ftrs is None: lin_ftrs = [50]
     if ps is None:  ps = [0.1]
     ds = data.train_ds
     vocab_size, n_class = len(data.vocab.itos), data.c
-    layers = [emb_sz * 3] + lin_ftrs + [n_class]
+    if bicls_head == 'BiPoolingLinearClassifier':
+        count = 3*2
+    else:
+        count = 3
+    layers = [emb_sz * count] + lin_ftrs + [n_class]
     ps = [dps[4]] + ps
     model = get_birnn_classifier(bptt, max_len, n_class, vocab_size, emb_sz, nh, nl, pad_token,
                                layers, ps, input_p=dps[0], weight_p=dps[1], embed_p=dps[2], hidden_p=dps[3],
-                               qrnn=qrnn)
+                               qrnn=qrnn, bicls_head=bicls_head)
     learn = RNNLearner(data, model, bptt, split_func=birnn_classifier_split, **kwargs)
     return learn
 
diff --git a/fastai_contrib/models.py b/fastai_contrib/models.py
index 0c328a7..b5de84c 100644
--- a/fastai_contrib/models.py
+++ b/fastai_contrib/models.py
@@ -138,14 +138,18 @@ def get_bilm(vocab_sz:int, emb_sz:int, n_hid:int, n_layers:int, pad_token:int, t
 
 def get_birnn_classifier(bptt:int, max_seq:int, n_class:int, vocab_sz:int, emb_sz:int, n_hid:int, n_layers:int,
                        pad_token:int, layers:Collection[int], drops:Collection[float], bidir:bool=False, qrnn:bool=False,
-                       hidden_p:float=0.2, input_p:float=0.6, embed_p:float=0.1, weight_p:float=0.5)->nn.Module:
+                       hidden_p:float=0.2, input_p:float=0.6, embed_p:float=0.1, weight_p:float=0.5, bicls_head:str='BiPoolingLinearClassifier')->nn.Module:
     "Create a RNN classifier model."
     fwd_rnn_enc = MultiBatchRNNCore(bptt, max_seq, vocab_sz, emb_sz, n_hid, n_layers, pad_token=pad_token, bidir=bidir,
                       qrnn=qrnn, hidden_p=hidden_p, input_p=input_p, embed_p=embed_p, weight_p=weight_p)
     bwd_rnn_enc = MultiBatchRNNCore(bptt, max_seq, vocab_sz, emb_sz, n_hid, n_layers, pad_token=pad_token, bidir=bidir,
                                 qrnn=qrnn, hidden_p=hidden_p, input_p=input_p, embed_p=embed_p, weight_p=weight_p)
 
-    model = SequentialRNN(BiLMModel(fwd_rnn_enc, bwd_rnn_enc), BiPoolingLinearClassifier(layers, drops))
+    head = BiPoolingLinearClassifier
+    if bicls_head == 'BiPoolingLinearClassifier': head = BiPoolingLinearClassifier
+    elif bicls_head == 'AvgPoolingLinearClassifier': head = AvgPoolingLinearClassifier
+
+    model = SequentialRNN(BiLMModel(fwd_rnn_enc, bwd_rnn_enc), head(layers, drops))
     model.reset()
     return model
 
diff --git a/ulmfit/train_clas.py b/ulmfit/train_clas.py
index 5581498..139f8a4 100644
--- a/ulmfit/train_clas.py
+++ b/ulmfit/train_clas.py
@@ -44,6 +44,8 @@ class CLSHyperParams(LMHyperParams):
     # dir_path -> data/imdb/
     use_test_for_validation=False
 
+    bicls_head:str = 'BiPoolingLinearClassifier'
+
     def __post_init__(self, *args, **kwargs):
         super().__post_init__(*args, **kwargs)
         self.dataset_dir=self.dataset_path
@@ -92,7 +94,10 @@ class CLSHyperParams(LMHyperParams):
         fastai.text.learner.default_dropout['language'] = dps or self.dps
         trn_args=dict(drop_mult=self.drop_mult, bptt=self.bptt, clip=self.clip,)
         trn_args.update(kwargs)
-        classifier_learner = bilm_text_classifier_learner if self.bidir else text_classifier_learner
+        classifier_learner = text_classifier_learner
+        if self.bidir:
+            classifier_learner = bilm_text_classifier_learner
+            trn_args['bicls_head'] = self.bicls_head
         learn = classifier_learner(data_clas,  pad_token=PAD_TOKEN_ID,
             path=self.model_dir.parent, model_dir=self.model_dir.name,
             qrnn=self.qrnn, emb_sz=self.emb_sz, nh=self.nh, nl=self.nl, **trn_args)
@@ -116,7 +121,6 @@ class CLSHyperParams(LMHyperParams):
         lm_val_df = lm_trn_df[:val_len]
 
         if use_test_for_validation:
-            val_len = max(int(len(tst_df) * 0.1), 2)
             val_df = tst_df
             cls_cache = 'notst'
         else:
@@ -126,8 +130,11 @@ class CLSHyperParams(LMHyperParams):
             cls_cache = '.'
 
         if self.tokenizer is Tokenizers.SUBWORD:
-            #TODO Fix me to make sure it trains correct dictionary
-            args = get_sentencepiece(self.dataset_path, self.dataset_path / 'train.csv', self.name, vocab_size=self.max_vocab)
+            args = get_sentencepiece(self.dataset_path, self.dataset_path / 'train.csv',
+                                     self.name, vocab_size=self.max_vocab, pre_rules=[], post_rules=[])
+            if self.tokenizer is Tokenizers.SUBWORD:
+                args = get_sentencepiece(self.dataset_path, self.dataset_path / 'train.csv',
+                                         self.name, vocab_size=self.max_vocab, pre_rules=[], post_rules=[])
         elif self.tokenizer is Tokenizers.MOSES:
             args = dict(tokenizer=Tokenizer(tok_func=MosesTokenizerFunc, lang='en', pre_rules=[], post_rules=[]))
         elif self.tokenizer is Tokenizers.MOSES_FA: