Fix assertion error

2020-07-06 09:23:53 -06:00 · 2020-07-06 09:23:53 -06:00 · 2636d3b620
commit 2636d3b620
parent 8f92c0a088
1 changed files with 1 additions and 1 deletions
--- a/codes/models/archs/SRG1_arch.py
+++ b/codes/models/archs/SRG1_arch.py
@ -66,7 +66,7 @@ class SwitchComputer(nn.Module):
        tc = transform_count
        if self.enable_negative_transforms:
            tc = transform_count * 2
-        assert filters > transform_count * 2
+        assert filters > tc
        self.final_switch_conv = nn.Conv2d(filters, tc, 1, 1, 0)

        self.transforms = nn.ModuleList([transform_block() for _ in range(transform_count)])