diff --git a/mhcflurry/class1_neural_network.py b/mhcflurry/class1_neural_network.py
index 735f0a7f9cefee344d618242eef412b6121370ea..9e6bc8b0a57d661228bfed1e3116229cf2625691 100644
--- a/mhcflurry/class1_neural_network.py
+++ b/mhcflurry/class1_neural_network.py
@@ -353,6 +353,13 @@ class Class1NeuralNetwork(object):
         result['prediction_cache'] = None
         return result
 
+    def __setstate__(self, state):
+        """
+        Deserialize. For pickle support.
+        """
+        self.__dict__.update(state)
+        self.prediction_cache = weakref.WeakKeyDictionary()
+
     def peptides_to_network_input(self, peptides):
         """
         Encode peptides to the fixed-length encoding expected by the neural
diff --git a/mhcflurry/select_allele_specific_models_command.py b/mhcflurry/select_allele_specific_models_command.py
index 722c82ab5154e5e15f3994ea5e8ddac223a3cfbe..18fb6e5196a01e7263d843d694171fb17cf32d03 100644
--- a/mhcflurry/select_allele_specific_models_command.py
+++ b/mhcflurry/select_allele_specific_models_command.py
@@ -153,6 +153,8 @@ def run(argv=sys.argv[1:]):
 
             df = df.loc[~df._excluded]
             print("Reduced data to: %s" % (str(df.shape)))
+
+        metadata_dfs["model_selection_data"] = df
     else:
         df = None
 
@@ -195,8 +197,7 @@ def run(argv=sys.argv[1:]):
         print("Attempting to create directory: %s" % args.out_models_dir)
         os.mkdir(args.out_models_dir)
         print("Done.")
-
-    metadata_dfs["model_selection_data"] = df
+    
     result_predictor = Class1AffinityPredictor(metadata_dataframes=metadata_dfs)
 
     worker_pool = worker_pool_with_gpu_assignments_from_args(args)