Fix per atom MAE bug

[EricZQu]

Broadcasting Bug in per_atom_mae Metrics

Affected Code

• src/fairchem/core/units/mlip_unit/_metrics.py — per_atom_mae, per_atom_mse
• src/fairchem/core/modules/evaluator.py — per_atom_mae, per_atom_mse

Energy Head Output Shape

The UMA energy head (escn_md.py, MLP_Energy_Head / MLP_EFS_Head) produces a 1-D energy tensor:

energy_part = torch.zeros(len(data_dict["natoms"]), ...)   # (S,)
energy_part.index_add_(0, data_dict["batch"], node_energy.view(-1))
return {"energy": energy}                                  # (S,)

Metric Inputs

compute_metrics masks predictions and targets but does not reshape them:

target_masked = batch[task.name][output_mask]   # (S',)
pred_masked   = pred[output_mask]               # (S',)
natoms_masked = batch.natoms[output_mask]       # (S',)

All three are 1-D with length S' (number of valid samples).

The Bug

def per_atom_mae(prediction, target, key):
    error = torch.abs(target[key] - prediction[key])   # (S',) - (S',) → (S',)
    return error / target["natoms"].unsqueeze(1)        # (S',) / (S', 1) → (S', S')  !!

The unsqueeze(1) converts natoms from (S',) to (S', 1). Dividing a 1-D tensor (S',) by a 2-D tensor (S', 1) triggers PyTorch broadcasting:

error:    (S',)    →  broadcast to (S', S')
natoms:   (S', 1)  →  broadcast to (S', S')
result:   (S', S')                             — should be (S',)

The @metrics_dict wrapper then computes mean, sum, and numel over the S' x S' matrix instead of a length-S' vector, producing wrong metric values.

per_atom_mse has the same bug.

Fix

Remove the unsqueeze(1):

def per_atom_mae(prediction, target, key):
    return torch.abs(target[key] - prediction[key]) / target["natoms"]