Fix lightning training

[zhiqwang]

No description provided.

[codecov]

Codecov Report

Merging #48 (d2241e9) into master (f34b032) will increase coverage by 6.06%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master      #48      +/-   ##
==========================================
+ Coverage   72.14%   78.21%   +6.06%     
==========================================
  Files           8        8              
  Lines         621      615       -6     
==========================================
+ Hits          448      481      +33     
+ Misses        173      134      -39     

Flag	Coverage Δ
unittests	78.21% <100.00%> (+6.06%)	⬆️

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

Impacted Files	Coverage Δ
test/dataset_utils.py	100.00% <100.00%> (+72.72%)	⬆️
test/test_engine.py	97.82% <100.00%> (+15.21%)	⬆️

---

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 f34b032...d2241e9. Read the comment docs.