Mean of norm2 unstable?

autograd
#1

I use a mean of norm2 difference as one of the losses the code is as follows:

def gradient_loss(input_depth, groundtruth_depth):
	gx_diff = (input_depth[:,:,:-1,:-1] - input_depth[:,:,1:,:-1]) - (groundtruth_depth[:,:,:-1,:-1] - groundtruth_depth[:,:,1:,:-1])
	gy_diff = (input_depth[:,:,:-1,:-1] - input_depth[:,:,:-1,1:]) - (groundtruth_depth[:,:,:-1,:-1] - groundtruth_depth[:,:,:-1,1:])
	diff_map = torch.cat((gx_diff, gy_diff), 1)
	diff_length = torch.norm(diff_map, 2, 1)
	return torch.mean(diff_length)

but get nan after one step optimization. When I comment this loss, nan is gone.

I am wondering whether this cost is unstable? Actually, I think this is a commonly used loss.

The whole loss is defined as:

import torch
import torch.nn.functional as F

def down_sample(depth_image):
	return depth_image[:,:,::2,::2]

def gradient_loss(input_depth, groundtruth_depth):
	gx_diff = (input_depth[:,:,:-1,:-1] - input_depth[:,:,1:,:-1]) - (groundtruth_depth[:,:,:-1,:-1] - groundtruth_depth[:,:,1:,:-1])
	gy_diff = (input_depth[:,:,:-1,:-1] - input_depth[:,:,:-1,1:]) - (groundtruth_depth[:,:,:-1,:-1] - groundtruth_depth[:,:,:-1,1:])
	diff_map = torch.cat((gx_diff, gy_diff), 1)
	diff_length = torch.norm(diff_map, 2, 1)
	return torch.mean(diff_length)

def build_loss(predict_depth, groundtruth_depth, gradient_weight = 1.0):
	# clamp depth_image between 0~50.0m and get inverse depth
	depth_image = groundtruth_depth.clamp(0.0, 50.0)
	depth_image = 1.0 / depth_image

	# get multi resolution depth maps
	depth_image1 = depth_image
	depth_image2 = down_sample(depth_image1)
	depth_image3 = down_sample(depth_image2)
	depth_image4 = down_sample(depth_image3)

	# build depth image loss
	loss_depth = F.l1_loss(predict_depth[0], depth_image1)
	loss_depth = loss_depth + F.l1_loss(predict_depth[1], depth_image2)
	loss_depth = loss_depth + F.l1_loss(predict_depth[2], depth_image3)
	loss_depth = loss_depth + F.l1_loss(predict_depth[3], depth_image4)

	# get the inv depth gradient
	loss_gradient = gradient_loss(predict_depth[0], depth_image1)
	loss_gradient = loss_gradient + gradient_loss(predict_depth[1], depth_image2)
	loss_gradient = loss_gradient + gradient_loss(predict_depth[2], depth_image3)
	loss_gradient = loss_gradient + gradient_loss(predict_depth[3], depth_image4)
	return loss_depth + gradient_weight * loss_gradient
#2

Now I figure it out that the norm do not right then input is 0. Use the master branch at hit may fix the problem.