about data_loader about snn_angular_velocity HOT 4 CLOSED

class RMSELoss:
    def __init__(self):
        self.str_to_idx = {
            'x': 0,
            'y': 1,
            'z': 2,
        }

    def compute(self, input_: torch.Tensor, target: torch.Tensor, time_start_idx: int=None, axis: str=None):
        assert len(input_.shape) == 3
        assert len(target.shape) == 3
        assert input_.shape == target.shape
        assert input_.shape[1] == 3
        """
        :param input: tensor of shape (batch, 3, time)
        :param target: tensor of shape (batch, 3, time)
        :param time_start_idx: Time-index from which to start computing the loss
        :param axis: {'x', 'y', 'z'}
        :return: Root mean squared error
        """
        if axis:
            axis_idx = self.str_to_idx[axis]
            input_ = input_[:, axis_idx, :]
            target = target[:, axis_idx, :]

        if time_start_idx:
            input_ = input_[..., time_start_idx:]
            target = target[..., time_start_idx:]

        return torch.mean(torch.sqrt(torch.mean((input_ - target) ** 2, dim=1)))