diff --git a/mhcflurry/antigen_presentation/presentation_component_models/presentation_component_model.py b/mhcflurry/antigen_presentation/presentation_component_models/presentation_component_model.py
index 308cde88c5a4e32fcbd5a794e395a72b6cc10238..52fe54a2fc93eb277d3c2b1e784fab827af382ab 100644
--- a/mhcflurry/antigen_presentation/presentation_component_models/presentation_component_model.py
+++ b/mhcflurry/antigen_presentation/presentation_component_models/presentation_component_model.py
@@ -233,7 +233,8 @@ class PresentationComponentModel(object):
             return_value = result_df[columns]
         if self.cached_predictions is not None:
             self.cached_predictions[cache_key] = return_value
-        return return_value
+        return dict(
+            (col, return_value[col].values) for col in self.column_names())
 
     def clone(self):
         """
diff --git a/mhcflurry/antigen_presentation/presentation_model.py b/mhcflurry/antigen_presentation/presentation_model.py
index f64b5000bb21f438235e3754c6f7ce0586cc7d34..fe8664e9ca6b68dd52ca6c059c49f39ea0c6d414 100644
--- a/mhcflurry/antigen_presentation/presentation_model.py
+++ b/mhcflurry/antigen_presentation/presentation_model.py
@@ -188,7 +188,7 @@ class PresentationModel(object):
                         model_input_training_hits_df)
                     self.trained_component_models[-1].append((sub_model,))
                     predictions = sub_model.predict(hits_and_decoys_df)
-                    for (col, values) in predictions.iteritems():
+                    for (col, values) in predictions.items():
                         hits_and_decoys_df[col] = values
                 final_predictor = self.fit_final_predictor(hits_and_decoys_df)
                 self.presentation_models_predictors.append(final_predictor)