Use w2nn.load_model to use cudnn #369

nagadomi · Oct 28, 2020 · c9b0fe7 · c9b0fe7
1 parent 44503fb
commit c9b0fe7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train.lua b/train.lua
@@ -526,7 +526,7 @@ local function train()
    }
    local model
    if settings.resume:len() > 0 then
-      model = torch.load(settings.resume, "ascii")
+      model = w2nn.load_model(settings.resume, settings.backend == "cudnn", "ascii")
       adam_config.xEvalCount = math.floor((#train_x * settings.patches) / settings.batch_size) * settings.batch_size * settings.inner_epoch * (settings.resume_epoch - 1)
       print(string.format("set eval count = %d", adam_config.xEvalCount))
       if adam_config.xEvalCount > 0 then