jax.devices issue with the finetuning notebook

[contango00]

is the versioning correct in the yml? JAX 4.26?

with the DEVICES = np.array(jax.devices()).reshape([1, 1, 1]) statement in the finetuning notebook I am getting the following error

ValueError: cannot reshape array of size 2 into shape (1,1,1)

The following....

        num_devices = jax.local_device_count()
        print(f'device kind: {jax.local_devices()[0].device_kind}')
        print(f'num_devices: {num_devices}')
        jax.devices()

returns this....

device kind: NVIDIA GeForce RTX 4080 SUPER
num_devices: 2
[cuda(id=0), cuda(id=1)]