SIGFPE error when training

When running the trainer, I get the following error:

W0628 23:34:14.945000 22654383867712 torch/multiprocessing/spawn.py:146] Terminating process 253794 via signal SIGTERM
Traceback (most recent call last):
  File ".../trellis2/TRELLIS/train.py", line 143, in <module>
    mp.spawn(main, args=(cfg,), nprocs=cfg.num_gpus, join=True)
  File ".../lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 282, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method="spawn")
  File ".../lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 238, in start_processes
    while not context.join():
  File ".../trellis2/lib/python3.10/site-packages/torch/multiprocessing/spawn.py", line 170, in join
    raise ProcessExitedException(
torch.multiprocessing.spawn.ProcessExitedException: process 0 terminated with signal SIGFPE

Through debugging, I narrowed down the location of the error to Line 304 in structured_latent_vae_mesh_dec.py. For reference, the line is, reps.extend(self.models['decoder'](args['latents'])). Note that, in my case, self.models['decoder'] is an instance of ElasticSLatMeshDecoder while args['latents'] is a sparse tensor.

As a final note, I am testing training on a smaller dataset using the command

python train.py \
  --config configs/vae/slat_vae_dec_mesh_swin8_B_64l8_fp16.json \
  --output_dir outputs/slat_vae_dec_mesh_swin8_B_64l8_fp16_1node \
  --num_nodes 1 \
  --auto_retry 0 \
  --data_dir <path to data> \

Thank you in advance for any help.

Edit:

As a note, I am using L4 gpus rather than A100s (and cannot switch to A100s) and am using Pytorch 2.4.1.

Start with trellis/trainers/vae/structured_latent_vae_mesh_dec.py at line 304 and reproduce using train.py with configs/vae/slat_vae_dec_mesh_swin8_B_64l8_fp16.json. Check the ElasticSLatMeshDecoder call with sparse latents under PyTorch 2.4.1 on L4 GPUs, then verify that the smaller training command completes without SIGFPE.