Decorator marking a test that requires a bigger GPU (24GB) for execution. Some example pipelines: Flux, SD3, Cog, etc.
(test_case)
| 597 | |
| 598 | |
| 599 | def require_big_gpu_with_torch_cuda(test_case): |
| 600 | """ |
| 601 | Decorator marking a test that requires a bigger GPU (24GB) for execution. Some example pipelines: Flux, SD3, Cog, |
| 602 | etc. |
| 603 | """ |
| 604 | if not is_torch_available(): |
| 605 | return pytest.mark.skip(reason="test requires PyTorch")(test_case) |
| 606 | |
| 607 | import torch |
| 608 | |
| 609 | if not torch.cuda.is_available(): |
| 610 | return pytest.mark.skip(reason="test requires PyTorch CUDA")(test_case) |
| 611 | |
| 612 | device_properties = torch.cuda.get_device_properties(0) |
| 613 | total_memory = device_properties.total_memory / (1024**3) |
| 614 | return pytest.mark.skipif( |
| 615 | total_memory < BIG_GPU_MEMORY, reason=f"test requires a GPU with at least {BIG_GPU_MEMORY} GB memory" |
| 616 | )(test_case) |
| 617 | |
| 618 | |
| 619 | def require_big_accelerator(test_case): |
nothing calls this directly
no test coverage detected
searching dependent graphs…