The problem is this:
        self.mat = torch.Tensor([[cos(self.roll) * cos(self.yaw), -cos(self.yaw) * sin(self.roll), sin(self.yaw), self.dx],
                            [cos(self.pitch) * sin(self.roll) + cos(self.roll) * sin(self.pitch) * sin(self.yaw),

which is breaking autograd and making self.mat a matrix fixed at init.
What you should do is create mat during forward and use torch.stack - not torch.Tensor - to do so. Quite likely, the speedup of using matrices is not that large.
Best regards
Thomas