File size: 1,096 Bytes
5fa1a76 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
The following skip decorators are used to set the requirements of tests CPU/GPU/TPU-wise: require_torch - this test will run only under torch require_torch_gpu - as require_torch plus requires at least 1 GPU require_torch_multi_gpu - as require_torch plus requires at least 2 GPUs require_torch_non_multi_gpu - as require_torch plus requires 0 or 1 GPUs require_torch_up_to_2_gpus - as require_torch plus requires 0 or 1 or 2 GPUs require_torch_tpu - as require_torch plus requires at least 1 TPU Let's depict the GPU requirements in the following table: | n gpus | decorator | |--------+--------------------------------| | >= 0 | @require_torch | | >= 1 | @require_torch_gpu | | >= 2 | @require_torch_multi_gpu | | < 2 | @require_torch_non_multi_gpu | | < 3 | @require_torch_up_to_2_gpus | For example, here is a test that must be run only when there are 2 or more GPUs available and pytorch is installed: python no-style @require_torch_multi_gpu def test_example_with_multi_gpu(): If a test requires tensorflow use the require_tf decorator. |