Update config.py

config.py

@@ -7,7 +7,7 @@ Precision: float32
 import jax
 import jax.numpy as jnp
 from flax import nnx
-from diffuse.models import CondUNet2D
+from diffuse.neural_network import CondUNet2D
 
 # Model definition
 model = CondUNet2D(