From 72b2a09106fd9009f877a448369b9d9414130e41 Mon Sep 17 00:00:00 2001
From: Ryan <23580140+brianhou0208@users.noreply.github.com>
Date: Tue, 6 May 2025 00:56:36 +0800
Subject: [PATCH] fix final norm only apply at last indice

---
 timm/models/rdnet.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/timm/models/rdnet.py b/timm/models/rdnet.py
index 3c556e37..246030af 100644
--- a/timm/models/rdnet.py
+++ b/timm/models/rdnet.py
@@ -318,8 +318,11 @@ class RDNet(nn.Module):
             feat_idx += 1
             x = stage(x)
             if feat_idx in take_indices:
-                # NOTE not bothering to apply norm_pre when norm=True as almost no models have it enabled
-                intermediates.append(x)
+                if norm and feat_idx == last_idx:
+                    x_inter = self.norm_pre(x)  # applying final norm to last intermediate
+                else:
+                    x_inter = x
+                intermediates.append(x_inter)
 
         if intermediates_only:
             return intermediates