| local base = import 'templates/base.libsonnet'; | |
| local tpus = import 'templates/tpus.libsonnet'; | |
| local utils = import "templates/utils.libsonnet"; | |
| local tputests = base.BaseTest { | |
| frameworkPrefix: 'kornia', | |
| modelName: 'tpu-tests', | |
| mode: 'postsubmit', | |
| configMaps: [], | |
| timeout: 900, # 15 minutes, in seconds. | |
| image: std.extVar('image'), | |
| imageTag: std.extVar('image-tag'), | |
| tpuSettings+: { | |
| softwareVersion: 'pytorch-VERSION', | |
| }, | |
| accelerator: tpus.v3_8, | |
| command: utils.scriptCommand( | |
| ||| | |
| python -c "import torch; print(torch.__version__)" | |
| python -c "import torch_xla; print(torch_xla.__version__)" | |
| python -c "import kornia; print(kornia.__version__)" | |
| pytest -v kornia/test/color kornia/test/enhance kornia/test/filters kornia/test/test_losses.py --device tpu --dtype float32 -k "not grad" | |
| test_exit_code=$? | |
| echo "\nFinished running commands.\n" | |
| test $test_exit_code -eq 0 | |
| ||| | |
| ), | |
| }; | |
| tputests.oneshotJob | |