diff --git a/tabfm/src/jax/model.py b/tabfm/src/jax/model.py
index 42f3019..7c6a7e4 100644
--- a/tabfm/src/jax/model.py
+++ b/tabfm/src/jax/model.py
@@ -32,7 +32,6 @@
 Y: Number of layers (a.k.a blocks)
 """
 
-from absl import logging
 import chex
 import math
 from math import pi