Support pytorch lightning 1.2.0