From e6ff657d4a063bfb08351c08fc3e11eb55d0c319 Mon Sep 17 00:00:00 2001 From: Marten Lienen Date: Sun, 2 May 2021 12:43:52 +0200 Subject: [PATCH] Restore the previous autograd setting Indiscriminately enabling autograd at the end of the loop also enables it when the user had explicitly disabled it before. This is a common occurrence when the loss is computed over a validation set. --- geomloss/sinkhorn_divergence.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/geomloss/sinkhorn_divergence.py b/geomloss/sinkhorn_divergence.py index 01f7f13..4875be3 100644 --- a/geomloss/sinkhorn_divergence.py +++ b/geomloss/sinkhorn_divergence.py @@ -175,6 +175,7 @@ def sinkhorn_loop( C_xxs, C_yys = [C_xxs], [C_yys] C_xys, C_yxs = [C_xys], [C_yxs] + prev_autograd = torch.is_grad_enabled() torch.autograd.set_grad_enabled(False) k = 0 # Scale index; we start at the coarsest resolution available @@ -219,7 +220,7 @@ def sinkhorn_loop( C_xy_, C_yx_ = C_xys[k + 1], C_yxs[k + 1] last_extrapolation = False # No need to re-extrapolate after the loop - torch.autograd.set_grad_enabled(True) + torch.autograd.set_grad_enabled(prev_autograd) else: # It's worth investing some time on kernel truncation... @@ -277,7 +278,7 @@ def sinkhorn_loop( C_xx, C_yy = C_xx_, C_yy_ C_xy, C_yx = C_xy_, C_yx_ - torch.autograd.set_grad_enabled(True) + torch.autograd.set_grad_enabled(prev_autograd) if last_extrapolation: # Last extrapolation, to get the correct gradients: