Skip to content

Commit d2345ef

Browse files
committed
Update test parameters for model size precision in test_model_summary.py
1 parent 4dc1dbc commit d2345ef

File tree

1 file changed

+7
-4
lines changed

1 file changed

+7
-4
lines changed

tests/tests_pytorch/utilities/test_model_summary.py

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -319,13 +319,16 @@ def test_empty_model_size(max_depth):
319319

320320

321321
@pytest.mark.parametrize(
322-
"accelerator",
322+
("accelerator", "precision"),
323323
[
324-
pytest.param("gpu", marks=RunIf(min_cuda_gpus=1)),
325-
pytest.param("mps", marks=RunIf(mps=True)),
324+
pytest.param("gpu", "16-true", marks=RunIf(min_cuda_gpus=1)),
325+
pytest.param("gpu", "32-true", marks=RunIf(min_cuda_gpus=1)),
326+
pytest.param("gpu", "64-true", marks=RunIf(min_cuda_gpus=1)),
327+
pytest.param("mps", "16-true", marks=RunIf(mps=True)),
328+
pytest.param("mps", "32-true", marks=RunIf(mps=True)),
329+
# Note: "64-true" with "mps" is skipped because MPS does not support float64
326330
],
327331
)
328-
@pytest.mark.parametrize("precision", ["16-true", "32-true", "64-true"])
329332
def test_model_size_precision(tmp_path, accelerator, precision):
330333
"""Test model size for different precision types."""
331334
model = PreCalculatedModel(precision=int(precision.split("-")[0]))

0 commit comments

Comments
 (0)