Formatting

Signed-off-by: Kunjan Patel <kunjan@ucla.edu>

Author: coolkp

src/maxdiffusion/checkpointing/checkpointing_utils.py

@@ -21,7 +21,6 @@
 import jax
 import numpy as np
 import os
-from jaxtyping import PyTree
 import orbax.checkpoint
 from maxdiffusion import max_logging
 from etils import epath