
  Note that we are doing weighted loss that only sum up over non-zero entries.

  Args:
    input_tensor: input tensor
    target_tensor: target tensor
    weight: weight tensor, in this case 1/sigma^2

  Returns:
    the weighted MSE loss
  """
  observation_dim = input_tensor.size()[-1]
  streched_tensor = ((input_tensor - target_tensor) ** 2).view(
      -1, observation_dim)