import torch
import numpy as np

@torch.no_grad()
def compute_iou(pred, target, num_classes):
    """Compute mean IoU."""
    pred = torch.argmax(pred, dim=1)
    iou_list = []
    for cls in range(num_classes):
        pred_mask = (pred == cls)
        target_mask = (target == cls)
        intersection = (pred_mask & target_mask).sum().float()
        union = (pred_mask | target_mask).sum().float()
        if union > 0:
            iou_list.append((intersection / union).item())
    return np.mean(iou_list) if iou_list else 0.0