Merge pull request fastai#171 from cbarokk/master

allow to pass a stepper class to fit()
dineshkumarsarangapani · Feb 23, 2018 · 34ea07e · 34ea07e
2 parents 1e74dde + a79dd2b
commit 34ea07e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/fastai/model.py b/fastai/model.py
@@ -61,7 +61,7 @@ def set_train_mode(m):
     else: m.train()
 
 
-def fit(model, data, epochs, opt, crit, metrics=None, callbacks=None, **kwargs):
+def fit(model, data, epochs, opt, crit, metrics=None, callbacks=None, stepper=Stepper, **kwargs):
     """ Fits a model
 
     Arguments:
@@ -72,7 +72,7 @@ def fit(model, data, epochs, opt, crit, metrics=None, callbacks=None, **kwargs):
        epochs(int): number of epochs
        crit: loss function to optimize. Example: F.cross_entropy
     """
-    stepper = Stepper(model, opt, crit, **kwargs)
+    stepper = stepper(model, opt, crit, **kwargs)
     metrics = metrics or []
     callbacks = callbacks or []
     avg_mom=0.98