Solved by ptrblck in post #2 
              
              
                Recreating a tensor (in rotational_matrix) will detach self.angle from the computation graph, so you should use torch.cat or torch.stack instead.