From f6ad788f733badf2377ea9f594a31d8f1f5fc024 Mon Sep 17 00:00:00 2001
From: luis <luis.guerrafernandez@monash.edu>
Date: Tue, 3 May 2022 20:06:26 -0600
Subject: [PATCH] I observed GPU memory leakage after every epoch, I couldn't
 find the root cause but emptying the cache fixed it.

---
 train.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/train.py b/train.py
index 610fbcf..893948a 100644
--- a/train.py
+++ b/train.py
@@ -619,6 +619,8 @@ def main():
             if args.distributed and hasattr(loader_train.sampler, 'set_epoch'):
                 loader_train.sampler.set_epoch(epoch)
 
+            torch.cuda.empty_cache()
+
             train_metrics = train_one_epoch(
                 epoch, model, loader_train, optimizer, train_loss_fn, args,
                 lr_scheduler=lr_scheduler, saver=saver, output_dir=output_dir,