瀏覽代碼

Use w2nn.load_model to use cudnn #369

nagadomi 4 年之前
父節點
當前提交
c9b0fe7f41
共有 1 個文件被更改,包括 1 次插入1 次删除
  1. 1 1
      train.lua

+ 1 - 1
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