test of the network

[iacopo97]

Thanks for your work!! I have a problem, after the completion of a training in Navier Stokes case if I launch the test I will get an l2 loss NaN. Which version of jax have you used?

[sifanexisted]

Our JAX version is 0.5.0, but I don’t think this issue is related to the JAX version.

I just double-checked by running the code, and it works fine.

CUDA_VISIBLE_DEVICES=0 python3 main.py --config=configs/cvit_8x8.py --config.mode=eval

Please note that the eval only work for single GPU. Can you provide more details?

[iacopo97]

I will try with 0.5.0, Can you share your checkpoints?

[sifanexisted]

It's quite large—around a GB —so we can't easily share it. As a simple test, you can simply try commenting out the checkpoint restore code and running it to see if you still get the NaN L2 error.

Or do you think the issue only happens when using your checkpoint?

[iacopo97]

The issue happens also if I am not restoring the checkpoint in the evaluation

[sifanexisted]

Interesting… Here’s the log output when I run the code without restoring from any checkpoint.

/scratch/sifanw/anaconda3/lib/python3.11/site-packages/absl/flags/_validators.py:254: UserWarning: Flag --config has a non-None default value; therefore, mark_flag_as_required will pass even if flag is not specified in the command line!
  mark_flag_as_required(flag_name, flag_values)
INFO:2025-07-21 10:19:34,857:jax._src.xla_bridge:945: Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
I0721 10:19:34.857594 140476058246016 xla_bridge.py:945] Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
INFO:2025-07-21 10:19:34,859:jax._src.xla_bridge:945: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object file: No such file or directory
I0721 10:19:34.859215 140476058246016 xla_bridge.py:945] Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object file: No such file or directory
I0721 10:19:52.746432 140476058246016 checkpoint_manager.py:623] [process=0][thread=MainThread] CheckpointManager init: checkpointers=None, item_names=None, item_handlers=None, handler_registry=None
I0721 10:19:52.746718 140476058246016 composite_checkpoint_handler.py:508] Initialized registry DefaultCheckpointHandlerRegistry({('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonSaveArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7fc0c444aed0>, ('metrics', <class 'orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonRestoreArgs'>): <orbax.checkpoint._src.handlers.json_checkpoint_handler.JsonCheckpointHandler object at 0x7fc0c444aed0>}).
I0721 10:19:52.747066 140476058246016 abstract_checkpointer.py:35] orbax-checkpoint version: 0.11.2
I0721 10:19:52.747142 140476058246016 async_checkpointer.py:88] [process=0][thread=MainThread] Using barrier_sync_fn: <function get_barrier_sync_fn.<locals>.<lambda> at 0x7fc205415300> timeout: 600 secs and primary_host=0 for async checkpoint writes
I0721 10:19:52.747678 140476058246016 checkpoint_manager.py:1535] Found 0 checkpoint steps in /scratch/sifanw/new_cvit/cvit/ns/ns_cvit_8x8/ckpt
I0721 10:19:52.747753 140476058246016 checkpoint_manager.py:1576] Saving root metadata
I0721 10:19:52.747825 140476058246016 checkpoint_manager.py:798] [process=0][thread=MainThread] CheckpointManager created,  primary_host=0, CheckpointManagerOptions=CheckpointManagerOptions(save_interval_steps=1, max_to_keep=10, keep_time_interval=None, keep_period=None, should_keep_fn=None, best_fn=None, best_mode='max', keep_checkpoints_without_metrics=True, step_prefix=None, step_format_fixed_length=None, step_name_format=None, create=True, cleanup_tmp_directories=False, save_on_steps=frozenset(), single_host_load_and_broadcast=False, todelete_subdir=None, enable_background_delete=False, read_only=False, enable_async_checkpointing=True, async_options=None, multiprocessing_options=MultiprocessingOptions(primary_host=0, active_processes=None, barrier_sync_key_prefix=None), should_save_fn=None, file_options=FileOptions(path_permission_mode=None), save_root_metadata=True, temporary_path_class=None, save_decision_policy=None), root_directory=/scratch/sifanw/new_cvit/cvit/ns/ns_cvit_8x8/ckpt: <orbax.checkpoint.checkpoint_manager.CheckpointManager object at 0x7fc0c47e4f90>
Total number of parameters: 92,135,171
/scratch/sifanw/anaconda3/lib/python3.11/multiprocessing/popen_fork.py:66: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.
  self.pid = os.fork()
l2_error: 2.9049873

[iacopo97]

Maybe the problem is related to the data? I am using this one: https://pdearena.github.io/pdearena/datadownload/

[iacopo97]

I am using jax version 0.6.2

[sifanexisted]

Yes, that’s possible. As far as I know, the dataset authors have updated it several times.

I’d suggest stepping through the code line by line to check whether the NaNs are coming from the model predictions or from the fetched evaluation mini-batches.

[iacopo97]

Can you share your dataset version?

[sifanexisted]

Unfortunately, I’m not sure how to check the exact dataset version. What I do know is that the total size is around 114 GB.

[iacopo97]

Ok, the NaN happens in the prediction phase, the coords and x values are not Nan

[iacopo97]

My dataset is around 81 GB

[iacopo97]

I use the version jax=0.5.0 and it works, i think that the jax version is really important