I have a total_ loss which is sum of -

A BCELoss
A Crossentropy loss
A custom loss function for image gradient.

The problem I am facing is that after 1st batch, some weights are updated to nan which results in all outputs as nan. If I remove the gradient loss, then it works fine.
What I found out was the denominator in the gradient loss were becoming 0, which was causing the problem. To fix it, I replace all denominators( Gradient magnitude) which were 0 with 1. But doing this only computes numerical loss for first backprop, next results in nan.
code-
sobel_x = torch.tensor([[+1, 0, -1], [+2, 0, -2], [+1, 0, -1]], requires_grad=False,dtype = torch.float)
	sobel_y = torch.tensor([[+1, +2, +1], [0, 0, 0], [-1, -2, -1]], requires_grad=False,dtype = torch.float)
	if cuda:
		sobel_x,sobel_y = sobel_x.cuda(),sobel_y.cuda()
		boundary_mask = boundary_mask.cuda()
	sobel_x = sobel_x.view((1,1,3,3))
	sobel_y = sobel_y.view((1,1,3,3))
	
	
	#gradients in the x and y direction for both predictions and the target transparencies
	G_x_pred = F.conv2d(pred,sobel_x,padding = 1)
	G_y_pred = F.conv2d(pred,sobel_y,padding = 1)
	G_x_target = F.conv2d(target,sobel_x,padding = 1)
	G_y_target = F.conv2d(target,sobel_y,padding = 1)

	#magnitudes of the gradients
	M_pred = torch.sqrt(torch.pow(G_x_pred,2)+torch.pow(G_y_pred,2))
	M_target = torch.sqrt(torch.pow(G_x_target,2)+torch.pow(G_y_target,2))
	
	#taking care of nans
	M_pred = (M_pred==0.).float() + M_pred
	M_target = (M_target==0.).float() + M_target

	# Lcos = (1-v_pred*v_target)*Magnitude_pred
	Lcos = (1-torch.abs((G_x_pred/M_pred)*(G_x_target/M_target)+(G_y_pred/M_pred)*(G_y_target/M_target)))*M_pred
	
	#Lmag = max(lambda*M_target-M_pred,0)
	lambd = 1.5
	Lmag = lambd*M_target-M_pred
	Lmag[Lmag<0] = 0

	gamma_1 = 0.5
	gamma_2 = 0.5

	#total gradient loss
	image_gradient_loss = (gamma_1*Lcos+gamma_2*Lmag)*boundary_mask

Can someone please help in tackling this situation.