fix

e3nn · Jun 24, 2023 · ab2a223 · ab2a223
1 parent 9d7f205
commit ab2a223
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/e3nn_jax/_src/tensor_products.py b/e3nn_jax/_src/tensor_products.py
@@ -14,6 +14,9 @@ def _prepare_inputs(input1, input2):
     input2 = e3nn.as_irreps_array(input2)
 
     dtype = get_pytree_dtype(input1, input2)
+    if dtype.kind == "i":
+        dtype = jnp.ones(()).dtype
+
     input1 = input1.astype(dtype)
     input2 = input2.astype(dtype)