diff --git a/fastai_contrib/utils.py b/fastai_contrib/utils.py
index 77d44a3..9004132 100644
--- a/fastai_contrib/utils.py
+++ b/fastai_contrib/utils.py
@@ -11,7 +11,6 @@ import re
 import csv
 
 from functools import reduce
-from fastai.text.data import TextDataset
 from fastai.text.transform import Tokenizer, BaseTokenizer, Vocab, default_rules
 from fastai.torch_core import *
 
@@ -80,7 +79,7 @@ def get_sentencepiece(path:PathOrStr, trn_path:Path, name:str, rules:ListRules=N
         sp_params = f'--input={raw_text_path} --pad_id={pad_idx} --unk_id=0' \
                     f'--character_coverage=1.0 --bos_id=-1 --eos_id=-1 ' \
                     f'--input_sentence_size={int(input_sentence_size)} ' \
-                    f'--model_prefix={path / 'models' / 'spm'} ' \
+                    f"--model_prefix={path / 'models' / 'spm'} " \
                     f'--vocab_size={vocab_size} --model_type={model_type} '
         spm.SentencePieceTrainer.Train(sp_params)
   
diff --git a/tests/test_end_to_end.py b/tests/test_end_to_end.py
index 77c77be..008c488 100644
--- a/tests/test_end_to_end.py
+++ b/tests/test_end_to_end.py
@@ -23,11 +23,13 @@ def check_data_exists():
 def test_pretrain_lm():
     imdb,wt2 = check_data_exists()
     lm_name="end-to-end-test-quick"
+    cuda_id=0
     results = ulmfit.pretrain_lm.pretrain_lm(
         dir_path=wt2,
         lang='en',
+        cuda_id=cuda_id,
         qrnn=True,
-        clean=True,
+        subword=False,
         max_vocab=1000,
         bs=80,
         num_epochs=1,
@@ -40,6 +42,7 @@ def test_pretrain_lm():
                 data_dir=get_data_folder(),
                 lang='en', pretrain_name=lm_name, model_dir=wt2/'models',
                 qrnn=True,
+                cuda_id=cuda_id,
                 fine_tune=True,
                 max_vocab=1000,
                 bs=20, bptt=70, name=lm_name+'-imdb-clas',
diff --git a/ulmfit/pretrain_lm.py b/ulmfit/pretrain_lm.py
index 26dfa53..0847ca6 100644
--- a/ulmfit/pretrain_lm.py
+++ b/ulmfit/pretrain_lm.py
@@ -56,9 +56,9 @@ def pretrain_lm(dir_path, lang='en', cuda_id=0, qrnn=True, subword=False, max_vo
     if qrnn:
         print('Using QRNNs...')
 
-    trn_path = dir_path / f'{lang}.wiki.train.tokens.unk'
-    val_path = dir_path / f'{lang}.wiki.valid.tokens.unk'
-    tst_path = dir_path / f'{lang}.wiki.test.tokens.unk'
+    trn_path = dir_path / f'{lang}.wiki.train.tokens'
+    val_path = dir_path / f'{lang}.wiki.valid.tokens'
+    tst_path = dir_path / f'{lang}.wiki.test.tokens'
     for path_ in [trn_path, val_path, tst_path]:
         assert path_.exists(), f'Error: {path_} does not exist.'