code: def test_callback_batch_on_device(tmpdir, accelerator):
      """Test that the batch object sent to the on_*_batch_start/end hooks is on the right device."""
      batch_callback = BatchHookObserverCallback()
  
      model = BatchHookObserverModel()
      trainer = Trainer(
          default_root_dir=tmpdir,
          max_steps=1,
          limit_train_batches=1,
          limit_val_batches=1,
          limit_test_batches=1,
          limit_predict_batches=1,
          accelerator=accelerator,
          devices=1,
          callbacks=[batch_callback],
      )
      trainer.fit(model)
      trainer.validate(model)
      trainer.test(model)
      trainer.predict(model)
  
imports: import pytest
  from lightning.pytorch import Callback, Trainer
  from lightning.pytorch.demos.boring_classes import BoringModel
  from tests_pytorch.helpers.runif import RunIf
label_desc: Use image node_id 2 for calling a global function with example usage: test_callback_batch_on_device(tmpdir, accelerator) without return types