diff --git a/mhcflurry/class1_neural_network.py b/mhcflurry/class1_neural_network.py
index 9f22176db08e141fbfc12f059a7c650941923c24..277c683b8b23c3750daa82e75c2be09a21f75059 100644
--- a/mhcflurry/class1_neural_network.py
+++ b/mhcflurry/class1_neural_network.py
@@ -1416,10 +1416,4 @@ class Class1NeuralNetwork(object):
                 original_model.fit = \
                 original_model.fit_generator = throw
 
-        session = K.get_session()
-        (weights_variable,) = layer.non_trainable_weights
-        session.run(weights_variable.initializer)
-        session.run(tf.assign(weights_variable, reshaped))
-        #
-        # Previously had this but was getting sporadic not initialized errors
-        # layer.set_weights([reshaped])
+        layer.set_weights([reshaped])
diff --git a/mhcflurry/testing_utils.py b/mhcflurry/testing_utils.py
index f6cc079b375acbad60ddac2e7974ac72377b06fe..157a11af9f62891be7e61d64f1fae93b5463e830 100644
--- a/mhcflurry/testing_utils.py
+++ b/mhcflurry/testing_utils.py
@@ -1,7 +1,13 @@
+"""
+Utilities used in MHCflurry unit tests.
+"""
 from . import Class1NeuralNetwork
 
 
-def module_cleanup():
+def cleanup():
+    """
+    Clear tensorflow session and other process-wide resources.
+    """
     import keras.backend as K
     Class1NeuralNetwork.clear_model_cache()
     K.clear_session()
diff --git a/test/test_calibrate_percentile_ranks_command.py b/test/test_calibrate_percentile_ranks_command.py
index 4c97e80c04e5d47c6b3ff9872bce5c32d015f602..2dc44a0b018a74113b5314bff7ecc79486115fb3 100644
--- a/test/test_calibrate_percentile_ranks_command.py
+++ b/test/test_calibrate_percentile_ranks_command.py
@@ -14,8 +14,8 @@ from mhcflurry.downloads import get_path
 
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 
-from mhcflurry.testing_utils import module_cleanup
-teardown = module_cleanup
+from mhcflurry.testing_utils import cleanup
+teardown = cleanup
 
 
 def run_and_check(n_jobs=0, delete=True, additional_args=[]):
diff --git a/test/test_changing_allele_representations.py b/test/test_changing_allele_representations.py
index 7c70048d0ed9ab6924cd184ac554a31a91cdae1f..38b9e12c339fd8bb2c6d40624c84bab92b4df636 100644
--- a/test/test_changing_allele_representations.py
+++ b/test/test_changing_allele_representations.py
@@ -8,8 +8,8 @@ from mhcflurry.downloads import get_path
 
 from numpy.testing import assert_equal
 
-from mhcflurry.testing_utils import module_cleanup
-teardown = module_cleanup
+from mhcflurry.testing_utils import cleanup
+teardown = cleanup
 
 ALLELE_TO_SEQUENCE = pandas.read_csv(
     get_path(
diff --git a/test/test_class1_affinity_predictor.py b/test/test_class1_affinity_predictor.py
index 563b6dcf3aa4830aa70b42a19a30896cec834773..8f2b4aa057bf1e42abc99752330ac45d28dded72 100644
--- a/test/test_class1_affinity_predictor.py
+++ b/test/test_class1_affinity_predictor.py
@@ -29,7 +29,7 @@ def setup():
 def teardown():
     global DOWNLOADED_PREDICTOR
     DOWNLOADED_PREDICTOR = None
-    mhcflurry.testing_utils.module_cleanup()
+    mhcflurry.testing_utils.cleanup()
 
 
 # To hunt down a weird warning we were seeing in pandas.
diff --git a/test/test_class1_neural_network.py b/test/test_class1_neural_network.py
index e614d602b576d4c4a9ecbeec14029eea8bcc1195..4620334a38e96b917f41ec9db515ac33384711ce 100644
--- a/test/test_class1_neural_network.py
+++ b/test/test_class1_neural_network.py
@@ -13,8 +13,8 @@ from mhcflurry.class1_neural_network import Class1NeuralNetwork
 from mhcflurry.downloads import get_path
 from mhcflurry.common import random_peptides
 
-from mhcflurry.testing_utils import module_cleanup
-teardown = module_cleanup
+from mhcflurry.testing_utils import cleanup
+teardown = cleanup
 
 
 def test_class1_neural_network_a0205_training_accuracy():
diff --git a/test/test_class1_pan.py b/test/test_class1_pan.py
index cdb483cf17260d93a59ba9ad940727393f9785bd..9ad87d69e1ced68490c29fca4a657ab2473da2a3 100644
--- a/test/test_class1_pan.py
+++ b/test/test_class1_pan.py
@@ -11,8 +11,8 @@ from mhcflurry import Class1AffinityPredictor,Class1NeuralNetwork
 from mhcflurry.allele_encoding import AlleleEncoding
 from mhcflurry.downloads import get_path
 
-from mhcflurry.testing_utils import module_cleanup
-teardown = module_cleanup
+from mhcflurry.testing_utils import cleanup
+teardown = cleanup
 
 
 HYPERPARAMETERS = {
diff --git a/test/test_custom_loss.py b/test/test_custom_loss.py
index 4c177666f378284ed99960224cd8376a297ed387..2426f48db03242e4427b1bf7fd59a527551401aa 100644
--- a/test/test_custom_loss.py
+++ b/test/test_custom_loss.py
@@ -11,8 +11,8 @@ import keras.backend as K
 
 from mhcflurry.custom_loss import CUSTOM_LOSSES
 
-from mhcflurry.testing_utils import module_cleanup
-teardown = module_cleanup
+from mhcflurry.testing_utils import cleanup
+teardown = cleanup
 
 
 def evaluate_loss(loss, y_true, y_pred):
diff --git a/test/test_download_models_class1.py b/test/test_download_models_class1.py
index 5e4332724c01ba845e80c4f13676f3c93385cce8..ffbf8b3cf6ac172cef426a2cadd9d0efa9ecd101 100644
--- a/test/test_download_models_class1.py
+++ b/test/test_download_models_class1.py
@@ -5,7 +5,7 @@ from numpy.testing import assert_equal
 
 from mhcflurry import Class1AffinityPredictor, Class1NeuralNetwork
 
-from mhcflurry.testing_utils import module_cleanup
+from mhcflurry.testing_utils import cleanup
 
 
 DOWNLOADED_PREDICTOR = None
@@ -19,7 +19,7 @@ def setup():
 def teardown():
     global DOWNLOADED_PREDICTOR
     DOWNLOADED_PREDICTOR = None
-    module_cleanup()
+    cleanup()
 
 
 def predict_and_check(
diff --git a/test/test_multi_output.py b/test/test_multi_output.py
index efbcccf0f9331239219133018e9043a548c2038f..b4254b56d7b1e7282c3127b00d15cf6593b98795 100644
--- a/test/test_multi_output.py
+++ b/test/test_multi_output.py
@@ -12,8 +12,8 @@ logging.getLogger('tensorflow').disabled = True
 from mhcflurry.class1_neural_network import Class1NeuralNetwork
 from mhcflurry.common import random_peptides
 
-from mhcflurry.testing_utils import module_cleanup
-teardown = module_cleanup
+from mhcflurry.testing_utils import cleanup
+teardown = cleanup
 
 
 def test_multi_output():
diff --git a/test/test_network_merging.py b/test/test_network_merging.py
index 8d1aca244fbdcc1572e696d704837cb76015c307..56683a52dbd8bfe63ad60b06ddaa76111d5d186f 100644
--- a/test/test_network_merging.py
+++ b/test/test_network_merging.py
@@ -5,7 +5,7 @@ from mhcflurry import Class1AffinityPredictor, Class1NeuralNetwork
 from mhcflurry.common import random_peptides
 from mhcflurry.downloads import get_path
 
-from mhcflurry.testing_utils import module_cleanup
+from mhcflurry.testing_utils import cleanup
 
 logging.getLogger('tensorflow').disabled = True
 
@@ -23,7 +23,7 @@ def setup():
 def teardown():
     global PAN_ALLELE_PREDICTOR
     PAN_ALLELE_PREDICTOR = None
-    module_cleanup()
+    cleanup()
 
 
 def test_merge():
diff --git a/test/test_predict_command.py b/test/test_predict_command.py
index 7d4d93406ae72715e6ac56e473c4178fb19f4a89..507d7c4a6c62939636ccb7315707957850a99c0a 100644
--- a/test/test_predict_command.py
+++ b/test/test_predict_command.py
@@ -6,8 +6,8 @@ from numpy.testing import assert_equal
 
 from mhcflurry import predict_command
 
-from mhcflurry.testing_utils import module_cleanup
-teardown = module_cleanup
+from mhcflurry.testing_utils import cleanup
+teardown = cleanup
 
 TEST_CSV = '''
 Allele,Peptide,Experiment
diff --git a/test/test_released_predictors_on_hpv_dataset.py b/test/test_released_predictors_on_hpv_dataset.py
index 02928f6db74a3a918be884f34e6311f9ff14e15f..3369257455157cc4047aedde6818139dee54b52f 100644
--- a/test/test_released_predictors_on_hpv_dataset.py
+++ b/test/test_released_predictors_on_hpv_dataset.py
@@ -13,7 +13,7 @@ from nose.tools import eq_, assert_less, assert_greater, assert_almost_equal
 from mhcflurry import Class1AffinityPredictor
 from mhcflurry.downloads import get_path
 
-from mhcflurry.testing_utils import module_cleanup
+from mhcflurry.testing_utils import cleanup
 
 
 def data_path(name):
@@ -41,7 +41,7 @@ def setup():
 def teardown():
     global PREDICTORS
     PREDICTORS = None
-    module_cleanup()
+    cleanup()
 
 
 def test_on_hpv(df=DF):
diff --git a/test/test_released_predictors_well_correlated.py b/test/test_released_predictors_well_correlated.py
index 67d0e98e266e86d4779c70569f9d166a17f35d3f..82071d653e3db2c6147e73d97e8d1f146b785d84 100644
--- a/test/test_released_predictors_well_correlated.py
+++ b/test/test_released_predictors_well_correlated.py
@@ -15,7 +15,7 @@ from mhcflurry.encodable_sequences import EncodableSequences
 from mhcflurry.downloads import get_path
 from mhcflurry.common import random_peptides
 
-from mhcflurry.testing_utils import module_cleanup
+from mhcflurry.testing_utils import cleanup
 
 PREDICTORS = None
 
@@ -33,7 +33,7 @@ def setup():
 def teardown():
     global PREDICTORS
     PREDICTORS = None
-    module_cleanup()
+    cleanup()
 
 
 def test_correlation(
diff --git a/test/test_speed.py b/test/test_speed.py
index 4f0d473e229b44160325ac7b184f53b16baa9157..eb604a5087458df6c517b7dde5db189d7ebe11e0 100644
--- a/test/test_speed.py
+++ b/test/test_speed.py
@@ -18,7 +18,7 @@ from mhcflurry.encodable_sequences import EncodableSequences
 from mhcflurry.common import random_peptides
 from mhcflurry.downloads import get_path
 
-from mhcflurry.testing_utils import module_cleanup
+from mhcflurry.testing_utils import cleanup
 
 ALLELE_SPECIFIC_PREDICTOR = Class1AffinityPredictor.load(
     get_path("models_class1", "models"))
@@ -43,7 +43,7 @@ def setup():
 def teardown():
     global PREDICTORS
     PREDICTORS = None
-    module_cleanup()
+    cleanup()
 
 
 DEFAULT_NUM_PREDICTIONS = 10000
diff --git a/test/test_train_and_related_commands.py b/test/test_train_and_related_commands.py
index 09ae60e569b9423fc92dc995777f23e8ffa5fde1..7b893f7be0914669115e0477284aa70fc300a88d 100644
--- a/test/test_train_and_related_commands.py
+++ b/test/test_train_and_related_commands.py
@@ -14,8 +14,8 @@ from numpy.testing import assert_array_less, assert_equal
 from mhcflurry import Class1AffinityPredictor
 from mhcflurry.downloads import get_path
 
-from mhcflurry.testing_utils import module_cleanup
-teardown = module_cleanup
+from mhcflurry.testing_utils import cleanup
+teardown = cleanup
 
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 
diff --git a/test/test_train_pan_allele_models_command.py b/test/test_train_pan_allele_models_command.py
index dd69e81053420e35a4588dfa848a8ef87190ab14..51c7c3ffb5bc16ba4e92504702e06fe8e4ee8d58 100644
--- a/test/test_train_pan_allele_models_command.py
+++ b/test/test_train_pan_allele_models_command.py
@@ -15,8 +15,8 @@ from numpy.testing import assert_equal, assert_array_less
 from mhcflurry import Class1AffinityPredictor,Class1NeuralNetwork
 from mhcflurry.downloads import get_path
 
-from mhcflurry.testing_utils import module_cleanup
-teardown = module_cleanup
+from mhcflurry.testing_utils import cleanup
+teardown = cleanup
 
 os.environ["CUDA_VISIBLE_DEVICES"] = ""