nan values when training student network (logits)

[nguyenvulong]

When running the logits, KD received nan values. I decided to modify the class below and realized that the student network didn't learn at all.
Update: This repo partially based on this. The original author also used log_softmax

class Logits(nn.Module):
	'''
	Do Deep Nets Really Need to be Deep?
	http://papers.nips.cc/paper/5484-do-deep-nets-really-need-to-be-deep.pdf
	'''
	def __init__(self):
		super(Logits, self).__init__()

	def forward(self, out_s, out_t):
		loss = F.mse_loss(out_s, out_t)
                #i printed the out_s here
                return loss

Figure: the tensors from out_s were vanishing quickly to nan

If anyone has the similar issue, please let me know.