diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py
index bc813cf..4fc0ceb 100644
--- a/trainer/trainer_utils.py
+++ b/trainer/trainer_utils.py
@@ -4,6 +4,7 @@
 import os
 import random
 import math
+import gc
 import numpy as np
 import torch
 import torch.distributed as dist