Print out energy information in model.py

[Frank-LIU-520]

I want to print out energy information at line 339-340 in module/models.py by running run_train.py with Mace model in line 245-255 in run_train.py. However, it doesn't print out these information. What should I do to get those information if I don't want to transmit those with return?

#line 245-255 in run_train.py
model = modules.ScaleShiftMACE( **model_config, correlation=args.correlation, gate=modules.gate_dict[args.gate], interaction_cls_first=modules.interaction_classes[ "RealAgnosticInteractionBlock" ], MLP_irreps=o3.Irreps(args.MLP_irreps), atomic_inter_scale=std, atomic_inter_shift=0.0, )

#Line 339-340 in models.py
` # Sum over nodes in graph
inter_e = scatter_sum(
src=node_inter_es, index=data["batch"], dim=-1, dim_size=num_graphs
) # [n_graphs,]

    # Add E_0 and (scaled) interaction energy
    total_energy = e0 + inter_e
    node_energy = node_e0 + node_inter_es
    **print(total_energy, node_energy)**`

[ilyes319]

You should use the logger with logging.info instead of print.