import added

Author: prishajain1

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -25,6 +25,7 @@
 from ..modeling_flax_utils import FlaxModelMixin, get_activation
 from ... import common_types
 from ..vae_flax import (FlaxAutoencoderKLOutput, FlaxDiagonalGaussianDistribution, FlaxDecoderOutput)
+import math
 
 BlockSizes = common_types.BlockSizes