def test_callback_batch_on_device(tmp_path, 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=tmp_path,
        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)
Test that the batch object sent to the on_*_batch_start/end hooks is on the right device.
tests/tests_pytorch/loops/test_all.py