Profiling PyTorch: pred.cpu() is reponsible for 95% of runtime?

[Marvin_Teichmann]

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…