Formatting

Signed-off-by: Kunjan <kunjanp@google.com>

Author: coolkp

src/maxdiffusion/train_utils.py

@@ -17,7 +17,6 @@
 import numpy as np
 import jax
 import jax.numpy as jnp
-import threading
 import queue
 
 from maxdiffusion import max_utils, max_logging

src/maxdiffusion/trainers/sdxl_trainer.py

@@ -24,7 +24,6 @@
 import jax.numpy as jnp
 from jax.sharding import PartitionSpec as P
 from flax.linen import partitioning as nn_partitioning
-from jax.experimental import io_callback
 from concurrent.futures import ThreadPoolExecutor
 from maxdiffusion.trainers.stable_diffusion_trainer import (StableDiffusionTrainer)