Updated pytorch_lightning import statements

[VikaasVarma]

With the newer versions of pytorch the import: from pytorch_lightning.utilities.distributed becomes from pytorch_lightning.utilities.rank_zero

[salesforce-cla[bot]]

Thanks for the contribution! Before we can merge this, we need @VikaasVarma to sign the Salesforce Inc. Contributor License Agreement.