Hi @KFrank, many thanks for your solution. The solution makes sense and it could definitely work for what I need to do. I am implementing a new type of classifier and one of the functions used in the square root of a positive semi-definite  matrix.
One quick question though, if one can do this elegantly like what you have proposed, then why are people working on a PyTorch function for this which involves more complicated math/algorithms, for eg. in this Github link below especially when you see the more recent comments from the last couple of days?

  
      github.com/pytorch/pytorch
  
  
    
  
	  
  

  
    
      [feature request] symmetric matrix square root 37
    

    
      
        opened 
        
          
        
        Aug 30, 2019
      
      


      
        
          
          yaroslavvb
        
      
    
  



  This is needed when incorporating curvature information into optimization
There's PyTorch implementation of symmetric matrix square root op here but they use...



    function request
    module: linear algebra
    triaged