Callback on_train_batch_end batch on CPU

[import-antigravity]

Hi, I have a callback with methods for on_train_batch_end and on_validation_batch_end, and I want to take the same batch and run it through another model. However, the data is on the CPU so I have to move it back to the GPU. Is there any way to avoid this redundant moving of data? It's really slowing things down.

[awaelchli]

Can you give a runnable example that demonstrates this? I can't see where Lightning would move these outputs back to CPU.

def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
    # pl_module, outputs, batch should all be on the same device

[import-antigravity]

Hopefully this works

[awaelchli]

Okay, I understand now.
This happens because we move the batch to the device in the accelerator right before it calls training_step(). The old batch reference is still on the CPU and this will be passed to the on_train_batch_end.

I will create an issue for this. Thanks for reporting.

[import-antigravity]

Happy to help!

[import-antigravity]

Okay, I understand now.
This happens because we move the batch to the device in the accelerator right before it calls training_step(). The old batch reference is still on the CPU and this will be passed to the on_train_batch_end.

I will create an issue for this. Thanks for reporting.

I don't know if you've created an issue yet but I made #6990 around the same time I made the OP

[awaelchli]

#6990 was probably a misunderstanding.

I created the issue here #7377 and linked a PR. Hopefully finishing it soon fixing the benchmark test.