From 4583add687d2d5a532b0b85fc5934bb135b56d87 Mon Sep 17 00:00:00 2001 From: Franck Mamalet Date: Wed, 5 Jun 2024 22:37:30 +0200 Subject: [PATCH] detach to avoid graph expansion --- deel/torchlip/modules/loss.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deel/torchlip/modules/loss.py b/deel/torchlip/modules/loss.py index 1492fa2..6520328 100644 --- a/deel/torchlip/modules/loss.py +++ b/deel/torchlip/modules/loss.py @@ -223,7 +223,7 @@ def _update_mean(self, y_pred): self.alpha_mean * self.current_mean + (1 - self.alpha_mean) * current_global_mean ) - self.current_mean = self.clamp_current_mean(current_global_mean) + self.current_mean = self.clamp_current_mean(current_global_mean).detach() total_mean = current_global_mean total_mean = torch.clamp(total_mean, self.min_margin_v, 20000) return total_mean