@@ -319,13 +319,16 @@ def test_empty_model_size(max_depth):
319
319
320
320
321
321
@pytest .mark .parametrize (
322
- "accelerator" ,
322
+ ( "accelerator" , "precision" ) ,
323
323
[
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
326
330
],
327
331
)
328
- @pytest .mark .parametrize ("precision" , ["16-true" , "32-true" , "64-true" ])
329
332
def test_model_size_precision (tmp_path , accelerator , precision ):
330
333
"""Test model size for different precision types."""
331
334
model = PreCalculatedModel (precision = int (precision .split ("-" )[0 ]))
0 commit comments