[FEATURE REQUEST] Implement loss weighing by class

[KristinaUlicna]

grace/grace/training/train.py

Lines 84 to 95 in bd91bdf

	# Specify node & edge criterion:
	# TODO: Implement class weighting
	node_criterion = torch.nn.CrossEntropyLoss(
	weight=None,
	reduction="mean",
	ignore_index=node_masked_class,
	)
	edge_criterion = torch.nn.CrossEntropyLoss(
	weight=None,
	reduction="mean",
	ignore_index=edge_masked_class,
	)