diff --git a/mhcflurry/custom_loss.py b/mhcflurry/custom_loss.py
index 8bd1b9537ce868f71f1a27bea5739fb98bad32cf..35440ca09fdf3462f4fae3ad2287524b9ce206f2 100644
--- a/mhcflurry/custom_loss.py
+++ b/mhcflurry/custom_loss.py
@@ -130,9 +130,6 @@ class MSEWithInequalitiesAndMultipleOutputs(object):
     def loss(y_true, y_pred):
         from keras import backend as K
 
-        #y_true = K.print_tensor(y_true, "y_true1")
-        #y_pred = K.print_tensor(y_pred, "y_pred1")
-
         y_true = K.flatten(y_true)
 
         output_indices = y_true // 10
@@ -143,8 +140,6 @@ class MSEWithInequalitiesAndMultipleOutputs(object):
         ordinals = K.arange(K.shape(y_true)[0])
         flattened_indices = (
             ordinals * y_pred.shape[1] + K.cast(output_indices, "int32"))
-        import tensorflow
-        #flattened_indices = tensorflow.Print(flattened_indices, [flattened_indices], "flattened_indices", summarize=1000)
         updated_y_pred = K.gather(K.flatten(y_pred), flattened_indices)
 
         # Alternative implementation using tensorflow, which could be used if
diff --git a/test/test_multi_output.py b/test/test_multi_output.py
index 31e75b0d7a4132881f4d4f0acaea42462f100e10..974316271574e1aa0ee6f06929bb69a95dc7ac1a 100644
--- a/test/test_multi_output.py
+++ b/test/test_multi_output.py
@@ -20,7 +20,7 @@ def test_multi_output():
         loss="custom:mse_with_inequalities_and_multiple_outputs",
         activation="tanh",
         layer_sizes=[16],
-        max_epochs=500,
+        max_epochs=50,
         minibatch_size=250,
         random_negative_rate=0.0,
         random_negative_constant=0.0,