Ok, you have one other problem as well: your compute_total_loss
function is not correct. Please see this thread with a checklist of potential issues. You’ve got one of the ones on the list. If the question is why the transposes are required, please have a look at this thread.