diff --git a/generator/ctrl/training_utils/training.py b/generator/ctrl/training_utils/training.py
index c31a224..a59965c 100644
--- a/generator/ctrl/training_utils/training.py
+++ b/generator/ctrl/training_utils/training.py
@@ -86,10 +86,10 @@ embedding_dim = 1280
 # we defer the transformer definition to transformer.py
 # here, we only define the tied softmax layer
 # this layer ties the softmax weights to the input embeddings
-with tf.device('/cpu:0'):
-    class TiedEmbeddingSoftmax(tf.keras.layers.Layer):
+class TiedEmbeddingSoftmax(tf.keras.layers.Layer):
 
-        def __init__(self, vocab_size=vocab_size, embedding_size=embedding_dim, **kwargs):
+    def __init__(self, vocab_size=vocab_size, embedding_size=embedding_dim, **kwargs):
+        with tf.device('/cpu:0'):
             super(TiedEmbeddingSoftmax, self).__init__()
             self.w = self.add_weight(name='w', shape=(vocab_size, embedding_size),
                                      initializer='random_normal',
@@ -98,7 +98,8 @@ with tf.device('/cpu:0'):
                                      initializer='zeros',
                                      trainable=True)
 
-        def call(self, inputs, embed=True):
+    def call(self, inputs, embed=True):
+        with tf.device('/cpu:0'):
             if embed:
                 dtype = tf.keras.backend.dtype(inputs)
                 if dtype != 'int32' and dtype != 'int64':
@@ -108,27 +109,25 @@ with tf.device('/cpu:0'):
                 return tf.tensordot(inputs, tf.transpose(self.w), 1) + self.b
 
 
-    # input for the keras model
-    tokens = tf.keras.layers.Input(shape=(seq_length,), dtype='int32')
+# input for the keras model
+tokens = tf.keras.layers.Input(shape=(seq_length,), dtype='int32')
 
 
-with tf.device('/cpu:0'):
-    # instantiates a tied softmax class
-    tied_embedding_softmax = TiedEmbeddingSoftmax()
+# instantiates a tied softmax class
+tied_embedding_softmax = TiedEmbeddingSoftmax()
 
-    # embedded tokens, before passing it to the transformer
-    embedded = tied_embedding_softmax(tokens, embed=True)
+# embedded tokens, before passing it to the transformer
+embedded = tied_embedding_softmax(tokens, embed=True)
 
 # the activations after passing it from the transformer
 # for some odd reason, TPUs don't play well with specifying the arguments of the Encoder() function
 # so you have to leave them at their defaults
 transformed = transformer.Encoder()(embedded, training=False)
 
-with tf.device('/cpu:0'):
-    # pass the activations from our tiedsoftmax class
-    # this time with embed=False denoting that we are doing the softmax operation
-    # and not a lookup
-    logits = tied_embedding_softmax(transformed, embed=False)
+# pass the activations from our tiedsoftmax class
+# this time with embed=False denoting that we are doing the softmax operation
+# and not a lookup
+logits = tied_embedding_softmax(transformed, embed=False)
 
 # finally, define the Keras model with inputs as tokens and outputs as the logits we just computed
 model = tf.keras.Model(inputs=tokens, outputs=logits)