Here is a pytorch geometric code that crashes when used with torchrun --nproc-per-node=2 torchrun.py import os
import time
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler

from edges import *
from regressors import *
from training import *
from mlflow_utils import *

from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts
from torch_geometric.data import Batch

import mlflow
mlflow.set_tracking_uri("file:../mlruns")
mlflow.set_experiment("C52-C138")

# === CONFIG ===
size = 1024
workers = 0
multigpu = True  # still needed if you want to run this also on a single GPU

configs = [
    {"name": "C52-100"},
    # {"name": "C104-120"},
    # {"name": "C132-138"},
    # {"name": "C102"},
    # {"name": "C122-128"},
]


def main():
    local_rank = int(os.environ["LOCAL_RANK"])
    rank = int(os.environ["RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu")

    if multigpu:
        torch.cuda.set_device(device)
        dist.init_process_group(backend="nccl", init_method="env://")
        if rank == 0:
            print(f"[DDP] Initialized process {rank} of {world_size}", flush=True)

    if rank == 0:
        mlflow.start_run(run_name="all-6")

    # === Dataset Processing ===
    train_set, test_set = process_and_combine_cached_datasets(
        configs, root="data", upcast=True, compressed=True
    )

    if multigpu:
        dist.barrier()

    set_deterministic(42)

    # === Data Loaders ===
    if multigpu:
        train_sampler = DistributedSampler(train_set, num_replicas=world_size, rank=rank, shuffle=True, seed=42)
        train_loader = DataLoader(train_set, batch_size=size, sampler=train_sampler, num_workers=workers, pin_memory=True)
    else:
        train_loader = DataLoader(train_set, batch_size=size, shuffle=True, num_workers=workers, pin_memory=True)

    test_loader = DataLoader(test_set, batch_size=size, num_workers=workers, pin_memory=True)

    # === Model ===
    model = TransformerRegressor(
        hidden_dim=12, edge_dim=1, pooling='set2set', mlp_factor=4,
        dropout=0.0, num_layers=6, heads=6, use_embeddings=True
    ).to(device)

    if multigpu:
        model = DDP(model, device_ids=[local_rank])

    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=2e-4,
        weight_decay=0.001,
        betas=(0.9, 0.999),
        eps=1e-8,
        fused=True
    )

    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode='min', factor=0.5, patience=8, min_lr=1e-6,
    )

    if rank == 0:
        describe_model(model)
        log_full_experiment_config(model=model, optimizer=optimizer, scheduler=scheduler)

    best_val_loss = float('inf')
    save_path = "best_model.pt"
    start_epoch = 1
    loss_type = 'mse'

    for epoch in range(start_epoch, 401):
        if multigpu:
            train_sampler.set_epoch(epoch)

        start_wall = time.time()
        start_cpu = time.process_time()

        train_loss = train(model, train_loader, optimizer, loss_fn, loss_type=loss_type)

        if rank == 0 and mlflow.active_run():
            mlflow.log_metric("train_loss_mse", train_loss, step=epoch)

        val_loss, loss_table = evaluate_and_log_per_size(
            model, test_loader, epoch=epoch, overall_loss_type=loss_type
        )

        if epoch == 150:
            for param_group in optimizer.param_groups:
                param_group['lr'] = 5e-5
            scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=60, T_mult=1, eta_min=2e-6)
        elif epoch == 210:
            scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=80, T_mult=1, eta_min=2e-6)

        scheduler.step(val_loss if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau) else None)

        wall_time = time.time() - start_wall
        cpu_time = time.process_time() - start_cpu
        current_lr = optimizer.param_groups[0]['lr']

        if rank == 0:
            print(f"Epoch {epoch:03d}: Train Loss = {train_loss:.6e}, Val Loss = {val_loss:.6e}, "
                  f"LR = {current_lr:.6f} Wall Time: {wall_time:.2f} s CPU Time: {cpu_time:.2f} s", flush=True)

        if val_loss < best_val_loss:
            best_val_loss = val_loss
            if rank == 0:
                torch.save(
                    model.module.state_dict() if multigpu else model.state_dict(),
                    save_path
                )
                save_predictions(model, test_loader, filepath="results_best.txt")
                plot_loss_per_size(epoch)
                print(f"Saved new best model at epoch {epoch} with val loss {val_loss:.6e}", flush=True)

        if rank == 0:
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.module.state_dict() if multigpu else model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'val_loss': best_val_loss,
            }, 'checkpoint.pt')

    if rank == 0:
        model.load_state_dict(torch.load(save_path, map_location=device))
        model.eval()
        save_predictions(model, test_loader, filepath="results_best.txt")

    if multigpu:
        dist.destroy_process_group()


if __name__ == "__main__":
    main(). Here is the full error: torchrun --nproc-per-node=2 torchrun.py 
W0710 22:56:50.140000 13375 site-packages/torch/distributed/run.py:766] 
W0710 22:56:50.140000 13375 site-packages/torch/distributed/run.py:766] *****************************************
W0710 22:56:50.140000 13375 site-packages/torch/distributed/run.py:766] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
W0710 22:56:50.140000 13375 site-packages/torch/distributed/run.py:766] *****************************************
[DDP] Initialized process 0 of 2
[LOAD] C52-100 train: data/processed/C52-100_train_77e2e446_flat.pt.zst
[LOAD] C52-100 train: data/processed/C52-100_train_77e2e446_flat.pt.zst
[LOAD] C52-100 test: data/processed/C52-100_test_a7940bfb_flat.pt.zst
[LOAD] C52-100 test: data/processed/C52-100_test_a7940bfb_flat.pt.zst
/venv/main/lib/python3.12/site-packages/torch/distributed/distributed_c10d.py:4631: UserWarning: No device id is provided via init_process_group or barrier . Using the current device set by the user. 
  warnings.warn(  # warn only once
[rank1]:[W710 22:57:23.126778537 ProcessGroupNCCL.cpp:4718] [PG ID 0 PG GUID 0 Rank 1]  using GPU 1 as device used by this process is currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect. You can pecify device_id in init_process_group() to force use of a particular device.
/venv/main/lib/python3.12/site-packages/torch/distributed/distributed_c10d.py:4631: UserWarning: No device id is provided via init_process_group or barrier . Using the current device set by the user. 
  warnings.warn(  # warn only once
[rank0]:[W710 22:57:24.064157983 ProcessGroupNCCL.cpp:4718] [PG ID 0 PG GUID 0 Rank 0]  using GPU 0 as device used by this process is currently unknown. This can potentially cause a hang if this rank to GPU mapping is incorrect. You can pecify device_id in init_process_group() to force use of a particular device.
W0710 22:57:26.523000 13375 site-packages/torch/distributed/elastic/multiprocessing/api.py:900] Sending process 13441 closing signal SIGTERM
E0710 22:57:27.190000 13375 site-packages/torch/distributed/elastic/multiprocessing/api.py:874] failed (exitcode: -11) local_rank: 0 (pid: 13440) of binary: /venv/main/bin/python3
Traceback (most recent call last):
  File "/venv/main/bin/torchrun", line 10, in <module>
    sys.exit(main())
             ^^^^^^
  File "/venv/main/lib/python3.12/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 355, in wrapper
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/venv/main/lib/python3.12/site-packages/torch/distributed/run.py", line 892, in main
    run(args)
  File "/venv/main/lib/python3.12/site-packages/torch/distributed/run.py", line 883, in run
    elastic_launch(
  File "/venv/main/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 139, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/venv/main/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 270, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
=======================================================
torchrun.py FAILED
-------------------------------------------------------
Failures:
  <NO_OTHER_FAILURES>
-------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2025-07-10_22:57:26
  host      : c3644b06480a
  rank      : 0 (local_rank: 0)
  exitcode  : -11 (pid: 13440)
  error_file: <N/A>
  traceback : Signal 11 (SIGSEGV) received by PID 13440
Help me finding the problem