diff --git a/comfy_extras/nodes_apg.py b/comfy_extras/nodes_apg.py
index 6e69b73f7..efa2a27c0 100644
--- a/comfy_extras/nodes_apg.py
+++ b/comfy_extras/nodes_apg.py
@@ -73,7 +73,7 @@ class APG(io.ComfyNode):
             guidance = cond - uncond
 
             if momentum != 0:
-                if not torch.is_tensor(running_avg):
+                if not torch.is_tensor(running_avg) or (running_avg.shape != guidance.shape):
                     running_avg = guidance
                 else:
                     running_avg = momentum * running_avg + guidance