Execute the decorated test only if a GPU is available. This function is intended to be applied to tests that require the precense of a CUDA GPU. If a CUDA GPU is absent, it will simply be skipped. Args: func: function to be annotated. If `func` is None, this method returns a decora
(func=None)
| 1351 | |
| 1352 | |
| 1353 | def run_cuda_only(func=None): |
| 1354 | """Execute the decorated test only if a GPU is available. |
| 1355 | |
| 1356 | This function is intended to be applied to tests that require the precense |
| 1357 | of a CUDA GPU. If a CUDA GPU is absent, it will simply be skipped. |
| 1358 | |
| 1359 | Args: |
| 1360 | func: function to be annotated. If `func` is None, this method returns a |
| 1361 | decorator the can be applied to a function. If `func` is not None this |
| 1362 | returns the decorator applied to `func`. |
| 1363 | |
| 1364 | Returns: |
| 1365 | Returns a decorator that will conditionally skip the decorated test method. |
| 1366 | """ |
| 1367 | |
| 1368 | def decorator(f): |
| 1369 | if tf_inspect.isclass(f): |
| 1370 | raise ValueError("`run_cuda_only` only supports test methods.") |
| 1371 | |
| 1372 | def decorated(self, *args, **kwargs): |
| 1373 | if not is_gpu_available(cuda_only=True): |
| 1374 | self.skipTest("Test requires CUDA GPU") |
| 1375 | |
| 1376 | return f(self, *args, **kwargs) |
| 1377 | |
| 1378 | return decorated |
| 1379 | |
| 1380 | if func is not None: |
| 1381 | return decorator(func) |
| 1382 | |
| 1383 | return decorator |
| 1384 | |
| 1385 | |
| 1386 | @tf_export("test.is_gpu_available") |
nothing calls this directly
no test coverage detected