provide an optional context in mx.load_checkpoint (#221)

dmlc · Mar 29, 2017 · 9474d95 · 9474d95
1 parent 1781290
commit 9474d95
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/src/model.jl b/src/model.jl
@@ -613,9 +613,14 @@ function load_checkpoint(prefix :: AbstractString, epoch :: Int)
   return (arch, arg_params, aux_params)
 end
 
-function load_checkpoint(prefix :: AbstractString, epoch :: Int, ::Type{FeedForward})
+"""
+    load_checkpoint(prefix, epoch, ::mx.FeedForward; context)
+
+Load a mx.FeedForward model from the checkpoint *prefix*, *epoch* and optionally provide a context.
+"""
+function load_checkpoint(prefix :: AbstractString, epoch :: Int, ::Type{FeedForward}; context = nothing)
   arch, arg_params, aux_params = load_checkpoint(prefix, epoch)
-  model = FeedForward(arch)
+  model = FeedForward(arch, context = context)
   model.arg_params = arg_params
   model.aux_params = aux_params
   return model