Rakefile in ruby-dnn-1.0.0 vs Rakefile in ruby-dnn-1.1.0

@@ -1,7 +1,9 @@
 require "bundler/gem_tasks"
 require "rake/testtask"
+require "yard"
+require "yard/rake/yardoc_task"
 
 Rake::TestTask.new(:test) do |t|
   t.libs << "test"
   t.libs << "ext"
   t.libs << "lib"
@@ -16,12 +18,14 @@
   sh "cd ext/rb_stb_image; make clean; unlink Makefile"
 end
 
 task :default => [:test, :build_rb_stb_image]
 
-task :doc do
-  src_list = Dir["lib/dnn.rb"]
-  src_list += Dir["lib/dnn/core/*.rb"]
-  src_list += Dir["lib/dnn/core/layers/*.rb"]
-  src_list += Dir["lib/dnn/*.rb"]
-  sh "yardoc #{src_list.join(' ')}"
+YARD::Rake::YardocTask.new do |t|
+  t.files = [
+    "lib/dnn.rb",
+    "lib/dnn/core/*.rb",
+    "lib/dnn/core/layers/*.rb",
+    "lib/dnn/*.rb",
+    "lib/dnn/datasets/*.rb",
+  ]
 end