diff --git a/examples/docs/example_2_wake_comparison.py b/examples/docs/example_2_wake_comparison.py
index 631f25d9f8b8132f251ea62ab20f2abc46490b15..5de212f4b2c86f26c936edd420023715f7061370 100644
--- a/examples/docs/example_2_wake_comparison.py
+++ b/examples/docs/example_2_wake_comparison.py
@@ -58,7 +58,8 @@ def main():
         def get_tf(windFarmModel):
             return TopFarmProblem(
                 design_vars=dict(zip('xy', init_pos.T)),
-                cost_comp=PyWakeAEPCostModelComponent(windFarmModel, n_wt=3, ws=10, wd=np.arange(0, 360, 12)),
+                cost_comp=PyWakeAEPCostModelComponent(windFarmModel, n_wt=3, ws=10, wd=np.arange(0, 360, 12),
+                                                      grad_method=None),
                 constraints=[SpacingConstraint(min_spacing),
                              XYBoundaryConstraint(boundary)],
                 driver=EasyScipyOptimizeDriver(),
diff --git a/topfarm/cost_models/py_wake_wrapper.py b/topfarm/cost_models/py_wake_wrapper.py
index 800dc0f3e4107d913bf00a6f62b65806c9d69f3c..9dab4842788a7ef7e59e766fcc81f739d3567480 100644
--- a/topfarm/cost_models/py_wake_wrapper.py
+++ b/topfarm/cost_models/py_wake_wrapper.py
@@ -13,6 +13,7 @@ from scipy.interpolate.interpolate import RegularGridInterpolator
 import warnings
 from py_wake.flow_map import HorizontalGrid
 import xarray as xr
+from py_wake.utils.gradients import autograd
 
 
 class PyWakeAEP(AEPCalculator):
@@ -55,17 +56,32 @@ class PyWakeAEP(AEPCalculator):
 
 
 class PyWakeAEPCostModelComponent(AEPCostModelComponent):
-    def __init__(self, windFarmModel, n_wt, wd=None, ws=None, max_eval=None, **kwargs):
+    def __init__(self, windFarmModel, n_wt, wd=None, ws=None, max_eval=None, grad_method=autograd, **kwargs):
         self.windFarmModel = windFarmModel
+
+        def aep(**kwargs):
+            return self.windFarmModel.aep(x=kwargs[topfarm.x_key],
+                                          y=kwargs[topfarm.y_key],
+                                          h=kwargs.get(topfarm.z_key, None),
+                                          type=kwargs.get(topfarm.type_key, 0),
+                                          wd=wd, ws=ws)
+
+        if grad_method:
+            dAEPdxy = self.windFarmModel.dAEPdxy(grad_method)
+
+            def daep(**kwargs):
+                return dAEPdxy(x=kwargs[topfarm.x_key],
+                               y=kwargs[topfarm.y_key],
+                               h=kwargs.get(topfarm.z_key, None),
+                               type=kwargs.get(topfarm.type_key, 0),
+                               wd=wd, ws=ws)
+        else:
+            daep = None
         AEPCostModelComponent.__init__(self,
                                        input_keys=[topfarm.x_key, topfarm.y_key],
                                        n_wt=n_wt,
-                                       cost_function=lambda **kwargs: self.windFarmModel.aep(
-                                           x=kwargs[topfarm.x_key],
-                                           y=kwargs[topfarm.y_key],
-                                           h=kwargs.get(topfarm.z_key, None),
-                                           type=kwargs.get(topfarm.type_key, 0),
-                                           wd=wd, ws=ws),
+                                       cost_function=aep,
+                                       cost_gradient_function=daep,
                                        output_unit='GWh',
                                        max_eval=max_eval, **kwargs)