Eliminate redundant device moves for callbacks

[import-antigravity]

🚀 Feature

Make it possible to access batches on callbacks without moving them to the CPU

Motivation

Currently, on callback hooks like on_train_batch_end where a batch tensor is available, the tensor is on cpu by default. If I want to, for example, run that batch through another model, I have to move it back to CUDA with batch.to('cuda') even though it must have been on CUDA in order for the training step to have taken place. This creates a serious performance bottleneck.

Pitch

Either add an argument to on_train_batch_end and similar hooks for cuda or add new hooks for before the move takes place.

Additional context

#6945

[tchaton]

Dear @import-antigravity,

I don't think it would be good practice to implement a second model forward in on_train_batch as Lightning won't be able to optimise properly.

You should encapsulate your logic in training_step directly and use manual_optimization if needed.

Best,
T.C

[import-antigravity]

To clarify, this is a second model I'd want to run the batch through in the callback, like a pre-trained network

[import-antigravity]

Turns out this is a bug, see #7377