rcan unshuffle fix

the-database · Jan 9, 2025 · bddaa7d · bddaa7d
1 parent 07e88d0
commit bddaa7d
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/tests/test_archs/test_archs.py b/tests/test_archs/test_archs.py
@@ -56,6 +56,11 @@
     {"upsampler": "geoensemblepixelshuffle"},
 ]
 
+EXTRA_ARCH_PARAMS["rcan"] = [
+    {"unshuffle_mod": True},
+    {"unshuffle_mod": False},
+]
+
 # A list of tuples in the format of (name, arch, scale, extra_params).
 FILTERED_REGISTRIES_SCALES_PARAMS = [
     (name, arch, scale, extra_params)

diff --git a/traiNNer/archs/rcan_arch.py b/traiNNer/archs/rcan_arch.py
@@ -277,7 +277,7 @@ def __init__(
             self.add_mean = nn.Identity()
 
         # define head module
-        unshuffle_mod = unshuffle_mod and scale < 4
+        unshuffle_mod = unshuffle_mod and scale <= 2
         self.downscale_factor = 1
         if unshuffle_mod:
             self.downscale_factor = 4 // scale