Resolve compatibility issues between optax and jax#1

Open

andizq wants to merge 1 commit intosmao-astro:masterfrom andizq:master

+3-3

Commits on Jan 18, 2024

Resolve compatibility issues between optax and jax leading to AttributeError: module 'jax.numpy' has no attribute 'DeviceArray', when real_time_prediction.py is run or called
andizq
committed