I’ve double-checked the size of the Tensors and they match the values shown in the comments. For reference, those sizes are
A:  torch.Size([1, 4, 4]) Detbar:  torch.Size([1])
U:  torch.Size([1, 4, 4])  S:  torch.Size([1, 4]) V:  torch.Size([1, 4, 4])
M:  torch.Size([1, 4, 4]) Cbar:  torch.Size([1, 4, 4])
diag_M:  torch.Size([1, 4])
rho_matrix:  torch.Size([1, 4, 4])
mask_off_diagonal:  torch.Size([1, 4, 4])
diag_Xi:  torch.Size([1, 4])
masked_presum_Xi:  torch.Size([1, 4, 4])
non_diag_Xi:  torch.Size([1, 4, 4])
Xi:  torch.Size([1, 4, 4])
Abar:  torch.Size([1, 4, 4])
DetBarBar:  torch.Size([1, 4, 4])



And in particular, the inputs for each multiplication op.

One line that could be an issue could be
diag_Xi = torch.sum((diag_M * rho_matrix)*mask_off_diagonal, dim=-1) #diag_Xi [1, N]

as it’s ([1, 4] * [1, 4, 4]) * [1, 4, 4]. What I wanted with this line here is to take the diagonal of M and element-wise multiply along each row of rho_matrix then zero out the diagonal elements then sum over the columns of each matrix within the batch (which in this case is a batch of 1).
or maybe
non_diag_Xi = (-M*rho_matrix)*(1 - torch.eye(M.shape[-1])) #non_diag_Xi = [1, N ,N]

as the 1 - torch.eye(M.shape[-1]) Tensor is [N, N].
I did reshape them to the same Tensor shape for element-wise multiplication but I got the same error.
Also, out of curiosity, how were you able to tell it’s a mul op in the stack trace that’s the error?
EDIT: One bit that’s confusing me is that I’m taking the gradient of the loss which is a scalar, and the first loss term contains the laplacian which only takes a single input (as that’s a restriction of torch.autograd.functional.hessian so why would it crash saying there’s a mismatch with size 25 and size 4? Surely it would only have size 1 and size 4? Unless, the crash is actually happen in the first backward call or the forward call where it would see the entire batch? Or, perhaps, could there be a mismatch in the forward method that results in the wrong dimension being summed over so the batch dimension remains?