From 13b380b86bd48885d33c862d3da30791246d56e6 Mon Sep 17 00:00:00 2001
From: "Jinwoo Park (Curt)" <www.jwpark.co.kr@gmail.com>
Date: Mon, 17 Mar 2025 15:45:21 +0900
Subject: [PATCH] Update functions.py

---
 schp/modules/functions.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/schp/modules/functions.py b/schp/modules/functions.py
index 7983730..4eff6ae 100644
--- a/schp/modules/functions.py
+++ b/schp/modules/functions.py
@@ -46,7 +46,7 @@ class InPlaceABN(autograd.Function):
         ctx.save_for_backward(x_norm, var, weight, bias)
         ctx.mark_non_differentiable(running_mean, running_var)
         
-        return x_norm
+        return x_norm, running_mean, running_var
     
     @staticmethod
     @once_differentiable