Fix the dtype selection

ziatdinovmax · Feb 4, 2025 · b5c7ac5 · b5c7ac5
1 parent 8024315
commit b5c7ac5
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/neurobayes/flax_nets/deterministic_nn.py b/neurobayes/flax_nets/deterministic_nn.py
@@ -50,8 +50,7 @@ def __init__(self,
         input_shape = (input_shape,) if isinstance(input_shape, int) else input_shape
         self.model = architecture
 
-        is_transformer = any(base.__name__.lower().find('transformer') >= 0 
-                        for base in architecture.__mro__)
+        is_transformer = 'transformer' in self.model.__class__.__name__.lower()
         input_dtype = jnp.int32 if is_transformer else jnp.float32
 
         if loss not in ['homoskedastic', 'heteroskedastic', 'classification']: