From 562d42fc0de86fd6edca66297d76999f47f356f4 Mon Sep 17 00:00:00 2001
From: Edgar Riba <edgar.riba@gmail.com>
Date: Sun, 9 Jan 2022 18:00:55 +0100
Subject: [PATCH] update to pytorch 1.10.1 (#1518)

* update to pytorch 1.10.1

* fix histogram tests and small refactor

* disable augmentation test

* enable more cases in test histogram

* Use of legacy backtick

* enable crop tests again
---
 .github/workflows/tests_cpu.yml |  2 +-
 kornia/enhance/histogram.py     | 27 ++++++++++++++++-----------
 setup_dev_env.sh                |  4 ++--
 test/enhance/test_histogram.py  |  4 +++-
 4 files changed, 22 insertions(+), 15 deletions(-)

diff --git a/.github/workflows/tests_cpu.yml b/.github/workflows/tests_cpu.yml
index 7fae7065cb..d2dba0ee1f 100644
--- a/.github/workflows/tests_cpu.yml
+++ b/.github/workflows/tests_cpu.yml
@@ -17,7 +17,7 @@ jobs:
       matrix:
         os: [ubuntu-latest, macos-latest, windows-latest]
         python-version: [3.6, 3.8, 3.9]
-        pytorch-version: [1.8.1, 1.9.1]
+        pytorch-version: [1.8.1, 1.9.1, 1.10.1]
     steps:
     - uses: actions/checkout@v2
     - name: Setup conda dependencies
diff --git a/kornia/enhance/histogram.py b/kornia/enhance/histogram.py
index 4088b37e5c..f71ef9c89f 100644
--- a/kornia/enhance/histogram.py
+++ b/kornia/enhance/histogram.py
@@ -214,25 +214,29 @@ def image_histogram2d(
 
     if bandwidth is None:
         bandwidth = (max - min) / n_bins
+
     if centers is None:
-        centers = min + bandwidth * (torch.arange(n_bins, device=image.device, dtype=image.dtype).float() + 0.5)
+        centers = min + bandwidth * (torch.arange(n_bins, device=image.device, dtype=image.dtype) + 0.5)
     centers = centers.reshape(-1, 1, 1, 1, 1)
+
     u = torch.abs(image.unsqueeze(0) - centers) / bandwidth
-    if kernel == "triangular":
-        mask = (u <= 1).to(u.dtype)
-        kernel_values = (1 - u) * mask
-    elif kernel == "gaussian":
+
+    if kernel == "gaussian":
         kernel_values = torch.exp(-0.5 * u ** 2)
-    elif kernel == "uniform":
+    elif kernel in ("triangular", "uniform", "epanechnikov",):
+        # compute the mask and cast to floating point
         mask = (u <= 1).to(u.dtype)
-        kernel_values = torch.ones_like(u, dtype=u.dtype, device=u.device) * mask
-    elif kernel == "epanechnikov":
-        mask = (u <= 1).to(u.dtype)
-        kernel_values = (1 - u ** 2) * mask
+        if kernel == "triangular":
+            kernel_values = (1. - u) * mask
+        elif kernel == "uniform":
+            kernel_values = torch.ones_like(u) * mask
+        else:  # kernel == "epanechnikov"
+            kernel_values = (1. - u ** 2) * mask
     else:
         raise ValueError(f"Kernel must be 'triangular', 'gaussian', " f"'uniform' or 'epanechnikov'. Got {kernel}.")
 
     hist = torch.sum(kernel_values, dim=(-2, -1)).permute(1, 2, 0)
+
     if return_pdf:
         normalization = torch.sum(hist, dim=-1, keepdim=True) + eps
         pdf = hist / normalization
@@ -248,4 +252,5 @@ def image_histogram2d(
         hist = hist.squeeze()
     elif image.dim() == 3:
         hist = hist.squeeze(0)
-    return hist, torch.zeros_like(hist, dtype=hist.dtype, device=hist.device)
+
+    return hist, torch.zeros_like(hist)
diff --git a/setup_dev_env.sh b/setup_dev_env.sh
index 5ac6691bac..7f5cd0b3a1 100755
--- a/setup_dev_env.sh
+++ b/setup_dev_env.sh
@@ -16,7 +16,7 @@ conda_bin=$conda_bin_dir/conda
 
 # download and install miniconda
 # check the operating system: Mac or Linux
-platform=`uname`
+platform=$(uname)
 if [[ "$platform" == "Darwin" ]];
 then
  download_link=https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-x86_64.sh
@@ -36,7 +36,7 @@ fi
 # define a python version to initialise the conda environment.
 # by default we assume python 3.7.
 python_version=${PYTHON_VERSION:-"3.7"}
-pytorch_version=${PYTORCH_VERSION:-"1.9.1"}
+pytorch_version=${PYTORCH_VERSION:-"1.10.1"}
 pytorch_mode=${PYTORCH_MODE:-""}  # use `cpuonly` for CPU or leave it in blank for GPU
 cuda_version=${CUDA_VERSION:-"10.2"}
 
diff --git a/test/enhance/test_histogram.py b/test/enhance/test_histogram.py
index dd3f67de8a..a1ddd6848b 100644
--- a/test/enhance/test_histogram.py
+++ b/test/enhance/test_histogram.py
@@ -51,7 +51,9 @@ def test_jit(self, device, dtype, kernel):
         op = TestImageHistogram2d.fcn
         op_script = torch.jit.script(op)
 
-        assert_close(op(*inputs), op_script(*inputs))
+        out, out_script = op(*inputs), op_script(*inputs)
+        assert_close(out[0], out_script[0])
+        assert_close(out[1], out_script[1])
 
     @pytest.mark.parametrize("kernel", ["triangular", "gaussian", "uniform", "epanechnikov"])
     @pytest.mark.parametrize("size", [(1, 1), (3, 1, 1), (8, 3, 1, 1)])