From 97b6354cd8ab52b250829f22edacf3e976b9f529 Mon Sep 17 00:00:00 2001 From: Sam-Armstrong Date: Wed, 24 Jul 2024 02:22:51 +0100 Subject: [PATCH] fix: update backends/jax __init__.py to reflect changes to module --- ivy/functional/backends/jax/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ivy/functional/backends/jax/__init__.py b/ivy/functional/backends/jax/__init__.py index 8468027de6618..bb7b9df6b381c 100644 --- a/ivy/functional/backends/jax/__init__.py +++ b/ivy/functional/backends/jax/__init__.py @@ -238,7 +238,7 @@ def closest_valid_dtype(type=None, /, as_native=False): from . import control_flow_ops from .control_flow_ops import * from . import module -from .module import Model +from .module import Module # sub-backends @@ -249,7 +249,7 @@ def closest_valid_dtype(type=None, /, as_native=False): if importlib.util.find_spec("flax"): import flax - NativeModule = Model + NativeModule = Module elif importlib.util.find_spec("haiku"): import haiku as hk