diff options
Diffstat (limited to 'tests/test_cuda_scaling.py')
| -rw-r--r-- | tests/test_cuda_scaling.py | 4 |
1 files changed, 3 insertions, 1 deletions
diff --git a/tests/test_cuda_scaling.py b/tests/test_cuda_scaling.py index 5f0befc..73cda66 100644 --- a/tests/test_cuda_scaling.py +++ b/tests/test_cuda_scaling.py @@ -15,9 +15,11 @@ def test_device_resolution(): def test_config_device_args(): - cfg = parse_args(["rnn", "0", "--device", "cpu", "--eval_batch_size", "256"]) + cfg = parse_args(["rnn", "0", "--device", "cpu", "--eval_batch_size", "256", "--use_amp", "True", "--compile_model", "True"]) assert cfg.device == "cpu" assert cfg.eval_batch_size == 256 + assert cfg.use_amp is True + assert cfg.compile_model is True def test_bisect_next_prime_large_range(): |