diff --git a/bilm/training.py b/bilm/training.py
index 2d60951..603292f 100644
--- a/bilm/training.py
+++ b/bilm/training.py
@@ -461,7 +461,7 @@ def _get_next_token_placeholders(suffix):
             # softmax_W is just the embedding layer
             self.softmax_W = self.embedding_weights
 
-        with tf.variable_scope('softmax'), tf.device('/cpu:0'):
+        with tf.variable_scope('softmax'), tf.device(lambda op: '' if not self.is_training else '/cpu:0'):
             # Glorit init (std=(1.0 / sqrt(fan_in))
             softmax_init = tf.random_normal_initializer(0.0,
                 1.0 / np.sqrt(softmax_dim))
@@ -502,10 +502,10 @@ def _get_next_token_placeholders(suffix):
 
                 else:
                     # get the full softmax loss
-                    output_scores = tf.matmul(
-                        lstm_output_flat,
-                        tf.transpose(self.softmax_W)
-                    ) + self.softmax_b
+                    output_scores = tf.transpose(tf.matmul(
+                        self.softmax_W, 
+                        tf.transpose(lstm_output_flat)
+                    )) + self.softmax_b
                     # NOTE: tf.nn.sparse_softmax_cross_entropy_with_logits
                     #   expects unnormalized output since it performs the
                     #   softmax internally