def test_amp_gpus(tmp_path, precision, devices):
    """Make sure combinations of AMP and strategies work if supported."""
    trainer = Trainer(
        default_root_dir=tmp_path,
        max_epochs=1,
        accelerator="gpu",
        devices=devices,
        strategy=("ddp_spawn" if devices > 1 else "auto"),
        precision=precision,
    )

    model = AMPTestModel()
    trainer.fit(model)
    trainer.test(model)
    trainer.predict(model, DataLoader(RandomDataset(32, 64)))
Make sure combinations of AMP and strategies work if supported.
tests/tests_pytorch/models/test_amp.py