Add parameter to specify which is the time dimension in MaskedMetric

[LucaButera]

tsl/tsl/metrics/torch/metric_base.py

Lines 112 to 122 in 154cccf

	def update(self, y_hat, y, mask=None):
	y_hat = y_hat[:, self.at]
	y = y[:, self.at]
	if mask is not None:
	mask = mask[:, self.at]
	if self.is_masked(mask):
	val, numel = self._compute_masked(y_hat, y, mask)
	else:
	val, numel = self._compute_std(y_hat, y)
	self.value += val
	self.numel += numel

In the referenced snippet, MaskedMetric's update function assumes the time dimension is the second.
However, this leads to adding an unnecessary dummy batch dimension when we represent a batch of graphs as a single big graph.

I suggest two possible solutions to avoid this:

1. Add a t_dim parameter and use x.select(t_dim, self.at) instead of x[:, self.at].
2. Add a pattern string and use it to identify the time dimension.

In both cases this can either be a class attribute or a method parameter, it depends on preferring to have it set once or allowing it to be changed each time the metric is updated.

The first solution is the easiest to implement, however the second one may make further dimension semantics dependent aggregations easier to implement down the road.

If this is deemed useful I can implement this behavior with an agreed solution.

[marshka]

Hi Luca, sorry for the late reply! Your suggestion seems useful and nice, I would opt for the first solution you proposed. Also, I think it is better to set the dim attribute just once at instantiation, at the moment I don't see any advantage in having it as a method parameter. Feel free to propose a PR!