Fix CV.

python-package/xgboost/callback.py

@@ -365,14 +365,22 @@ def before_training(self, model):
         '''Function called before training.'''
         for c in self.callbacks:
             model = c.before_training(model=model)
-            assert isinstance(model, Booster), 'before_training should return the Booster'
+            msg = 'before_training should return the model'
+            if self.is_cv:
+                assert isinstance(model.cvfolds, list), msg
+            else:
+                assert isinstance(model, Booster), msg
         return model
 
     def after_training(self, model):
         '''Function called after training.'''
         for c in self.callbacks:
             model = c.after_training(model=model)
-            assert isinstance(model, Booster), 'after_training should return the Booster'
+            msg = 'after_training should return the model'
+            if self.is_cv:
+                assert isinstance(model.cvfolds, list), msg
+            else:
+                assert isinstance(model, Booster), msg
         return model
 
     def before_iteration(self, model, epoch, dtrain, evals):

python-package/xgboost/training.py

@@ -493,9 +493,8 @@ def cv(params, dtrain, num_boost_round=10, nfold=3, stratified=False, folds=None
             verbose_eval, early_stopping_rounds, maximize, 0,
             num_boost_round, feval, None, callbacks,
             show_stdv=show_stdv, cvfolds=cvfolds)
-    callbacks.before_training(cvfolds)
-
     booster = _PackedBooster(cvfolds)
+    callbacks.before_training(booster)
 
     for i in range(num_boost_round):
         if callbacks.before_iteration(booster, i, dtrain, None):
@@ -522,4 +521,7 @@ def cv(params, dtrain, num_boost_round=10, nfold=3, stratified=False, folds=None
             results = pd.DataFrame.from_dict(results)
         except ImportError:
             pass
+
+    callbacks.after_training(booster)
+
     return results