Re-fix AQT conv construction.

PiperOrigin-RevId: 641333993
google-research · Jun 7, 2024 · 64d2f0e · 64d2f0e
1 parent 1ae4029
commit 64d2f0e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/chirp/models/efficientnet.py b/chirp/models/efficientnet.py
@@ -23,7 +23,7 @@
 from typing import Callable, NamedTuple
 
 from aqt.jax.v2 import aqt_conv_general
-from aqt.jax.v2 import aqt_dot_general
+from aqt.jax.v2 import config as aqt_cfg  # pylint: disable=unused-import
 from chirp.models import layers
 from flax import linen as nn
 import flax.typing as flax_typing
@@ -133,7 +133,7 @@ class OpSet:
         head_activation=nn.hard_swish,
         dot_general=jax.lax.dot_general,
         conv_general_dilated=aqt_conv_general.make_conv_general_dilated(
-            aqt_conv_general.conv_general_dilated_make(spatial_dimensions=2)
+            aqt_cfg.conv_general_dilated_make(spatial_dimensions=2)
         ),
     ),
 }