Fix lightning callback import. Add grad_clip_fn.

[rockerBOO]

Always errors if lightning is not available. The try method should handle this case.

• Added grad_clip_fn to add our own gradient clip function. Used with the accelerator library which interacts with the model, optimizer.

Maybe it would make more sense to have somewhere else than the step function because it wouldn't change probably.