diff --git a/beir/retrieval/train.py b/beir/retrieval/train.py
index f80f9f15..4592b17b 100644
--- a/beir/retrieval/train.py
+++ b/beir/retrieval/train.py
@@ -14,10 +14,9 @@
 )
 from sentence_transformers.readers import InputExample
 from torch import nn
-from torch.optim import Optimizer
+from torch.optim import Optimizer, AdamW
 from torch.utils.data import DataLoader
 from tqdm.autonotebook import trange
-from transformers import AdamW
 
 logger = logging.getLogger(__name__)