From 6c4e2b86418656d2be6e6e6bdd42b331604565fe Mon Sep 17 00:00:00 2001 From: Shubham Agarwal Date: Fri, 29 May 2020 00:06:44 +0100 Subject: [PATCH 1/2] cuda set device --- .gitignore | 2 ++ train.py | 1 + 2 files changed, 3 insertions(+) diff --git a/.gitignore b/.gitignore index 243a7b4..c997595 100644 --- a/.gitignore +++ b/.gitignore @@ -37,3 +37,5 @@ logs/ # virtualenv venv/ ENV/ + +.idea/ diff --git a/train.py b/train.py index 0e733ec..45a0520 100644 --- a/train.py +++ b/train.py @@ -113,6 +113,7 @@ if args.gpu_ids[0] >= 0 else torch.device("cpu") ) +torch.cuda.set_device(device) # Print config and args. print(yaml.dump(config, default_flow_style=False)) From 0e010a4bdc9200624cc68ba503f77eaff2043d50 Mon Sep 17 00:00:00 2001 From: Shubham Agarwal Date: Fri, 29 May 2020 00:08:50 +0100 Subject: [PATCH 2/2] empty cuda cache after epoch instead of batch --- train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/train.py b/train.py index 45a0520..5cb87f8 100644 --- a/train.py +++ b/train.py @@ -280,7 +280,7 @@ def lr_lambda_fun(current_iteration: int) -> float: scheduler.step(global_iteration_step) global_iteration_step += 1 - torch.cuda.empty_cache() + torch.cuda.empty_cache() # ------------------------------------------------------------------------- # ON EPOCH END (checkpointing and validation)