diff --git a/depthfm/unet/attention.py b/depthfm/unet/attention.py
index 6a30ab1..f2d2982 100644
--- a/depthfm/unet/attention.py
+++ b/depthfm/unet/attention.py
@@ -6,7 +6,7 @@ from inspect import isfunction
 import torch.nn.functional as F
 from typing import Optional, Any
 
-from util import checkpoint
+from .util import checkpoint
 
 
 try: