Aha! It was absolutely amazing! The way my output changed when I turned on anomaly detection on my own terminal! Thank you so much for the help, it got fixed  Apparently, the torch.stack was causing the problem, I changed it to
log_A_expanded = log_a.unsqueeze(2).expand((m, n, p))
log_B_expanded = log_b.unsqueeze(0).expand((m, n, p))

even though it basically works the same way but is not computing 
Though I am a little confused about why it might have failed in
	log_A_expanded = torch.stack([log_A] * p, dim=2)
	log_B_expanded = torch.stack([log_B] * m, dim=0)

Do you have any insight?