feat(ml): add lib rearrange

jolibrain · Aug 27, 2024 · 5eeafdf · 5eeafdf
1 parent 2167ebd
commit 5eeafdf
Showing 1 changed file with 1 addition and 0 deletions.
diff --git 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):