diff --git a/mhcflurry/class1_neural_network.py b/mhcflurry/class1_neural_network.py
index d75f66b71ddb3f330b4a1d3db76770d8a4014c3e..86a0cdf5a11766e02116e7c27ee9020a6b9fdfb3 100644
--- a/mhcflurry/class1_neural_network.py
+++ b/mhcflurry/class1_neural_network.py
@@ -1368,6 +1368,7 @@ class Class1NeuralNetwork(object):
                   m is the length of the vectors used to represent amino acids
         """
         from keras.models import clone_model
+        import keras.backend as K
         reshaped = allele_representations.reshape(
             (allele_representations.shape[0], -1))
         original_model = self.network()
diff --git a/test/test_class1_neural_network.py b/test/test_class1_neural_network.py
index 0566b6a610ce108c6bd68f2fb2e4672530ede772..8baaded0a02d88d43d6e77ad7551a687adc3528f 100644
--- a/test/test_class1_neural_network.py
+++ b/test/test_class1_neural_network.py
@@ -90,7 +90,7 @@ def test_inequalities():
     hyperparameters = dict(
         peptide_amino_acid_encoding="one-hot",
         activation="tanh",
-        layer_sizes=[16],
+        layer_sizes=[64],
         max_epochs=200,
         minibatch_size=32,
         random_negative_rate=0.0,
@@ -108,7 +108,7 @@ def test_inequalities():
         loss="custom:mse_with_inequalities_and_multiple_outputs")
 
     df = pandas.DataFrame()
-    df["peptide"] = random_peptides(1000, length=9)
+    df["peptide"] = random_peptides(100, length=9)
 
     # First half are binders
     df["binder"] = df.index < len(df) / 2
diff --git a/test/test_released_predictors_well_correlated.py b/test/test_released_predictors_well_correlated.py
index a6034b3ed8935d359f9203c15a29e2719dfa5c23..395a6e80b54c699c664d5b5ddbb9ff5990743648 100644
--- a/test/test_released_predictors_well_correlated.py
+++ b/test/test_released_predictors_well_correlated.py
@@ -23,7 +23,7 @@ PREDICTORS = {
         get_path("models_class1_pan", "models.with_mass_spec"))
 }
 
-PREDICTORS["pan-allele"].optimize()
+# PREDICTORS["pan-allele"].optimize()
 
 
 def test_correlation(