text: Title: Switch from pytorch_lightning to lightning
  Body: We should change all the imports from `pytorch_lightning` to `lightning`
  
  However, naively substituting all the occurrences does not work -- unknown reasons atm.