def test_dtensor_model_summary(_):
    """Test that the model summary can work with layers that have DTensor parameters."""
    # We mock the `_is_dtensor` to pretend parameters are DTensors, because testing with real DTensors
    # would require setting up distributed
    dtensor_model = UnorderedModel()
    summary = ModelSummary(dtensor_model)
    assert summary.total_layer_params > 0
    assert summary.total_parameters > 0
    assert summary.trainable_parameters > 0
Test that the model summary can work with layers that have DTensor parameters.
tests/tests_pytorch/utilities/test_model_summary.py