qwerfdsadad

I also encounter this problem. Have you solved this?

> The original project's multi-GPU training led to this problem Do you remember how to fix this bug? I tried many ways and still can't solve it.

> Reference Thanks for your reply. ```python vm_evolve = jax.pmap(jax.vmap(evolve, axis_name='j'), axis_name='i') local_devices = jax.local_device_count() uu = vm_evolve(u.reshape([local_devices, cfg.multi.numbers//local_devices, -1])) save_dim=[cfg.multi.numbers]+list(uu.shape[-2:]) uu_reshape=uu.reshape(save_dim) jnp.save(cwd+cfg.multi.save+'1D_Advection_Sols_beta'+str(beta)[:5], uu_reshape) ``` This is mt solution. For...

Have you successfully fixed this bug? For both FNO and Unet, I can't successfully run, and encountered the same problem as you. The dataset used comes from folder data_download. FNO...

@Mrzhang-dada @ionutmodo @osainz59 Have you solve this problem? I'm facing the same problem. the environment configuration shows blow: NVIDIA-SMI 535.129.03 Driver Version: 535.129.03 CUDA Version: 12.2 nvcc -V : 11.8...

I just use a simple model to test the functionality of deepseed, without using transformer.