diff --git a/torch_sgld/__init__.py b/torch_sgld/__init__.py
index 4b8c279..fdb011b 100644
--- a/torch_sgld/__init__.py
+++ b/torch_sgld/__init__.py
@@ -1,8 +1,8 @@
 from .sgld import SGLD
-from .lr_scheduler import CosineLR
+from .lr_scheduler import CosineLR, ABAnnealingLR
 
 __all__ = [
     'SGLD',
     'ABAnnealingLR',
     'CosineLR',
-]
\ No newline at end of file
+]