In the second stage (swa) with different weighting, the original code…

… used different loss function as compared with that chosed in the first stage. I modified the default loss function in the second stage as UnversalLoss, which may just be right for my purpose of studies. I would suggest to define a separate function and invoke it in these two stages to keep consistence.

mace/cli/run_train.py

@@ -620,12 +620,14 @@ def run(args: argparse.Namespace) -> None:
                 f"Using stochastic weight averaging (after {args.start_swa} epochs) with energy weight : {args.swa_energy_weight}, forces weight : {args.swa_forces_weight}, dipole weight : {args.swa_dipole_weight} and learning rate : {args.swa_lr}"
             )
         else:
-            loss_fn_energy = modules.WeightedEnergyForcesLoss(
+            loss_fn_energy = modules.UniversalLoss(
                 energy_weight=args.swa_energy_weight,
                 forces_weight=args.swa_forces_weight,
+                stress_weight=args.swa_stress_weight,
+                huber_delta=args.huber_delta,
             )
             logging.info(
-                f"Using stochastic weight averaging (after {args.start_swa} epochs) with energy weight : {args.swa_energy_weight}, forces weight : {args.swa_forces_weight} and learning rate : {args.swa_lr}"
+                f"Using stochastic weight averaging (after {args.start_swa} epochs) with energy weight : {args.swa_energy_weight}, forces weight : {args.swa_forces_weight}, stress weight : {args.swa_stress_weight} and learning rate : {args.swa_lr}"
             )
         swa = tools.SWAContainer(
             model=AveragedModel(model),