diff --git a/models/modules/diffusion_utils.py b/models/modules/diffusion_utils.py
index bf1ef9ff1..a364e96fb 100644
--- a/models/modules/diffusion_utils.py
+++ b/models/modules/diffusion_utils.py
@@ -2,6 +2,7 @@
 import torch
 import math
 from functools import partial
+from einops import rearrange
 
 
 def gamma_embedding_1D(gammas, dim, max_period):