diff --git a/connectomics/jax/checkpoint.py b/connectomics/jax/checkpoint.py
index 5ef121c..202ec9a 100644
--- a/connectomics/jax/checkpoint.py
+++ b/connectomics/jax/checkpoint.py
@@ -14,7 +14,6 @@
 # limitations under the License.
 """Utilities for model checkpointing."""
 
-import dataclasses
 import re
 from typing import Any, Optional, Sequence, TypeVar