I posted a follow up question using the insides gained by additional profiling runs in a     new thread.
Regarding forum etiquette. Is it generally preferred to post follow up questions in the same thread or to create a new thread for new questions?

  
    
    
    Applying torch.matmul along custom dimension 
  
  
    For an implementation of a graphical model I need to perform matmul along the axis 1 and 2 of a four-dimensional input tensor. The axis 0 and 3 should be broadcasted. The current implementation of torch matmul performs the matrix multiplication across the final two axis and performs broadcasting across all of input[:-2]. 
Currently I solve this by first transposing the input and then performing matmul. 
    input_row = input_col.transpose(1, 3)
    message = torch.matmul(self.gaussian, input_row…