torchcell.losses.mle_dist_supcr.MleDistSupCR

class MleDistSupCR(lambda_mse: float = 1.0, lambda_dist: float = 0.1, lambda_supcr: float = 0.001, dist_bandwidth: float = 0.5, supcr_temperature: float = 0.1, embedding_dim: int = 128, use_buffer: bool = True, buffer_size: int = 256, min_samples_for_dist: int = 64, min_samples_for_supcr: int = 64, use_ddp_gather: bool = True, gather_interval: int = 1, use_adaptive_weighting: bool = True, warmup_epochs: int | None = None, stable_epoch: int | None = None, use_temp_scheduling: bool = True, init_temperature: float = 1.0, final_temperature: float = 0.1, temp_schedule: str = 'exponential', weights: Tensor | None = None, max_epochs: int = 1000)

Bases: Module

Composite loss combining MSE, distribution matching, and SupCR.

Adds circular buffers, DDP synchronization, and adaptive weighting.

gather_across_gpus(tensor: Tensor) → Tensor

Gather tensors from all GPUs.

forward(predictions: Tensor, targets: Tensor, z_P: Tensor, epoch: int | None = None) → tuple[Tensor, dict[str, Any]]

Forward pass computing composite loss.

Parameters:

• predictions – Model predictions [batch_size, num_dims]

• targets – Target values [batch_size, num_dims]

• z_P – Perturbed embeddings for SupCR [batch_size, embedding_dim]

• epoch – Current training epoch (optional)

Returns:

Tuple of (total_loss, loss_dict)