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,