update

2020-09-05 22:00:17 +08:00
parent 39c754374c
commit e3c760d0c5
12 changed files with 122 additions and 43 deletions
--- a/main.py
+++ b/main.py
@@ -65,6 +65,8 @@ def run(task, config: str, *omega_options, **kwargs):
    backup_config = kwargs.get("backup_config", False)
    setup_output_dir = kwargs.get("setup_output_dir", False)
    setup_random_seed = kwargs.get("setup_random_seed", False)
+    assert torch.backends.cudnn.enabled
+    torch.backends.cudnn.benchmark = True
    with idist.Parallel(backend=backend) as parallel:
        parallel.run(running, conf, task, backup_config=backup_config, setup_output_dir=setup_output_dir,
                     setup_random_seed=setup_random_seed)