diff --git a/llava/train/llava_trainer.py b/llava/train/llava_trainer.py
index f8e0132bb..485292ad0 100755
--- a/llava/train/llava_trainer.py
+++ b/llava/train/llava_trainer.py
@@ -1,4 +1,6 @@
+import math
 import os
+import sys
 import torch
 import torch.nn as nn
 import datetime