Works like a charm! Thank you very much! That’s what I figured that it had gotten something to do with a broken graph…
Since I know the axis of the rotation and it is one of the three main axes of the coordinate system, the structure of the corresponding quaternion is already given. That’s why I don’t need to learn qx and qy, hence they can both be 0.
For the sake of completeness, the working final code of the nn.Module looks like this now:
class AngleModel(nn.Module):
    def __init__(self):
        super(AngleModel, self).__init__()
        self.angle = nn.Parameter(torch.tensor(0.0))

    def forward(self, input):
        qw = torch.cos(self.angle / 2.)
        qx = 0.0
        qy = 0.0
        qz = torch.sin(self.angle / 2.)

        matrix = torch.zeros(3, 3)

        matrix[0, 0] = 1. - 2. * qy ** 2 - 2. * qz ** 2
        matrix[1, 1] = 1. - 2. * qx ** 2 - 2. * qz ** 2
        matrix[2, 2] = 1. - 2. * qx ** 2 - 2. * qy ** 2

        matrix[0, 1] = 2. * qx * qy - 2. * qz * qw
        matrix[1, 0] = 2. * qx * qy + 2. * qz * qw

        matrix[0, 2] = 2. * qx * qz + 2 * qy * qw
        matrix[2, 0] = 2. * qx * qz - 2 * qy * qw

        matrix[1, 2] = 2. * qy * qz - 2. * qx * qw
        matrix[2, 1] = 2. * qy * qz + 2. * qx * qw

        output = torch.matmul(input, matrix)
        return output