I cannot run the training in NS case

[iacopo97]

I am using the same A6000 gpu that you have used but during training of cvit4x4 the computation explode.

[sifanexisted]

Thanks for asking! It might be due to a difference in the JAX version.

As a quick fix, try reducing the batch size and the number of query points in the config file.