diff --git a/tests/optimization_test.py b/tests/optimization_test.py
index dacec4099b..1675c8b496 100644
--- a/tests/optimization_test.py
+++ b/tests/optimization_test.py
@@ -20,7 +20,7 @@ import unittest
 
 import torch
 
-import optimization as optimization
+import optimization
 
 class OptimizationTest(unittest.TestCase):
 
diff --git a/tests/tokenization_test.py b/tests/tokenization_test.py
index 1cfc520ca1..8d6ede9300 100644
--- a/tests/tokenization_test.py
+++ b/tests/tokenization_test.py
@@ -19,7 +19,7 @@ from __future__ import print_function
 import os
 import unittest
 
-import tokenization as tokenization
+import tokenization
 
 
 class TokenizationTest(unittest.TestCase):