fix flux case

nod-ai · Feb 11, 2025 · 0654ef6 · 0654ef6
1 parent cdb80ef
commit 0654ef6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sharktank/sharktank/models/flux/testing.py b/sharktank/sharktank/models/flux/testing.py
@@ -11,8 +11,8 @@
 from .flux import FluxParams, FluxModelV1
 from .export import export_flux_transformer, flux_transformer_default_batch_sizes
 from ...types import DefaultPrimitiveTensor, Theta, save_load_theta
+from ...utils.testing import make_rand_torch
 from ...layers.testing import (
-    make_rand_torch,
     make_mmdit_double_block_random_theta,
     make_mmdit_single_block_random_theta,
 )