diff --git a/src/twomartens/masterthesis/aae/train_aae.py b/src/twomartens/masterthesis/aae/train_aae.py
index 13a1447..0733afb 100644
--- a/src/twomartens/masterthesis/aae/train_aae.py
+++ b/src/twomartens/masterthesis/aae/train_aae.py
@@ -46,9 +46,11 @@ from tensorflow.python.ops import summary_ops_v2
 
 from twomartens.masterthesis.aae import model
 from twomartens.masterthesis.aae import util
-from twomartens.masterthesis.aae.train import K
 from twomartens.masterthesis.aae.train import LOG_FREQUENCY
-from twomartens.masterthesis.aae.train import tfe
+
+# shortcuts for tensorflow sub packages and classes
+K = tf.keras.backend
+tfe = tf.contrib.eager
 
 GRACE: int = 10
 TOTAL_LOSS_GRACE_CAP: int = 6