diff --git a/common_dataset.py b/common_dataset.py
index dbed634aaaeed2d4aab3ac259ca286e7cd66c32e..b2fca949fbdde16e17c82c4ca3b6b9813cfcc37c 100644
--- a/common_dataset.py
+++ b/common_dataset.py
@@ -152,8 +152,8 @@ class Common_Dataset(Dataset):
 
 def load_data(path_train, path_val, path_test, batch_size, length, pad=False, convert=False, vocab = 'unmod'):
     print('Loading data')
-    data_train = pd.read_pickle(path_train)
     data_val = pd.read_pickle(path_val)
+    data_train = pd.read_pickle(path_train)
     data_test = pd.read_pickle(path_test)
     train = Common_Dataset(data_train, length, pad, convert, vocab)
     test = Common_Dataset(data_val, length, pad, convert, vocab)
diff --git a/model_custom.py b/model_custom.py
index 96cfc64035ba8c8c40d79b94b81246c05cc54005..eddc158c4f1dffcedaa3ae029222ab3862bc4193 100644
--- a/model_custom.py
+++ b/model_custom.py
@@ -1,7 +1,6 @@
 import math
 import torch.nn as nn
 import torch
-from tape import TAPETokenizer
 from tape.models.modeling_bert import ProteinBertModel
 
 class PermuteLayer(nn.Module):