train_GNN_stress.py