The foloowing multiGPU CUDA job is not working. It hangs: import os
import time
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
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 = 2
multigpu = True  # <--- set to False for single GPU

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


def main(rank, world_size):
    device = torch.device(f"cuda:{rank}")

    if multigpu:
        dist.init_process_group("nccl", rank=rank, world_size=world_size)
        torch.cuda.set_device(rank)
        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 ===
    if rank == 0:
        train_set, test_set = process_and_combine_cached_datasets(
            configs, root="data", upcast=True, compressed=True
        )
    if multigpu:
        if rank != 0:
            train_set, test_set = process_and_combine_cached_datasets(
                configs, root="data", upcast=True, compressed=True
            )
#        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=[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)

    # === Training ===
    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')

    # === Final Model Evaluation (rank 0 only) ===
    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()


# === DDP Launch ===
def run_training():
    if multigpu:
        world_size = torch.cuda.device_count()
        mp.spawn(main, args=(world_size,), nprocs=world_size, join=True)
    else:
        main(rank=0, world_size=1)

if __name__ == "__main__":
    run_training()
Help me finding the problem