diff --git a/tests/akro/test_discrete.py b/tests/akro/test_discrete.py
index 56c0407..1d4de11 100644
--- a/tests/akro/test_discrete.py
+++ b/tests/akro/test_discrete.py
@@ -3,9 +3,9 @@
 
 import numpy as np
 import pytest
-import tensorflow as tf
 
 from akro import Discrete
+from akro import tf
 from akro import theano
 from akro.requires import requires_tf, requires_theano
 
diff --git a/tests/akro/test_tuple.py b/tests/akro/test_tuple.py
index 9235e8e..f67b9ba 100644
--- a/tests/akro/test_tuple.py
+++ b/tests/akro/test_tuple.py
@@ -2,10 +2,10 @@
 import unittest
 
 import numpy as np
-import tensorflow as tf
 
 from akro import Box
 from akro import Discrete
+from akro import tf
 from akro import theano
 from akro import Tuple
 from akro.requires import requires_tf, requires_theano