Remove unused vmap in TracingWhile

EMI-Group · Jan 13, 2025 · f1cdcd1 · f1cdcd1
1 parent 4957c6d
commit f1cdcd1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/evox/utils/_control_flow/loop.py b/src/evox/utils/_control_flow/loop.py
@@ -4,7 +4,7 @@
 
 import torch
 
-from ...core import ModuleBase, _vmap_fix, jit, jit_class, trace_impl, use_state, vmap, vmap_impl
+from ...core import ModuleBase, _vmap_fix, jit, jit_class, trace_impl, use_state, vmap_impl
 from ...core.module import UseStateFunc
 from .utils import VarArgsCallable, VarArgsCallableMultiRet, _get_cache_key_object