Inherit Datasets from LightningDataModule

[zhiqwang]

This is a follow-up PR of #43 .

[codecov]

Codecov Report

Merging #46 (b29f94d) into master (3171d65) will not change coverage.
The diff coverage is n/a.

@@           Coverage Diff           @@
##           master      #46   +/-   ##
=======================================
  Coverage   78.35%   78.35%           
=======================================
  Files           8        8           
  Lines         619      619           
=======================================
  Hits          485      485           
  Misses        134      134           

Flag	Coverage Δ
unittests	78.35% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
test/test_engine.py	97.72% <ø> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 3171d65...b29f94d. Read the comment docs.