def test_model_saves_on_multi_gpu(tmp_path):
    """Test that ONNX model saves on a distributed backend."""
    trainer_options = {
        "default_root_dir": tmp_path,
        "max_epochs": 1,
        "limit_train_batches": 10,
        "limit_val_batches": 10,
        "accelerator": "gpu",
        "devices": [0, 1],
        "strategy": "ddp_spawn",
        "enable_progress_bar": False,
    }

    model = BoringModel()
    model.example_input_array = torch.randn(5, 32)

    tpipes.run_model_test(trainer_options, model, min_acc=0.08)

    file_path = os.path.join(tmp_path, "model.onnx")
    model.to_onnx(file_path)
    assert os.path.exists(file_path) is True
Test that ONNX model saves on a distributed backend.
tests/tests_pytorch/models/test_onnx.py