diff options
| author | CaptainJack2491 <jayrupnakawala@gmail.com> | 2026-08-17 16:35:28 +0100 |
|---|---|---|
| committer | CaptainJack2491 <jayrupnakawala@gmail.com> | 2026-08-17 16:35:28 +0100 |
| commit | 7166dc47c9bd764e2cd1d7c3d9e046cdc439b58c (patch) | |
| tree | 93159dcc7f6bb53726dac2c39ec7d4258c501b84 /tests/test_cuda_scaling.py | |
| parent | 66b8b77feb5fafd3696044ee33a80f26cb693a13 (diff) | |
perf(cuda): add AMP fp16 autocast with GradScaler and torch.compile reduce-overhead mode
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(): |