diff --git a/micro_sam/training/medsam_trainer.py b/micro_sam/training/medsam_trainer.py
index b05c3f55..b65dfc60 100644
--- a/micro_sam/training/medsam_trainer.py
+++ b/micro_sam/training/medsam_trainer.py
@@ -1,7 +1,9 @@
+import os
 import time
 
 import torch
 from torchvision.utils import make_grid
+
 import torch_em
 from torch_em.trainer.logger_base import TorchEmLogger