From 912a4d9feabf6040fd652f500f2b0b8f5077722a Mon Sep 17 00:00:00 2001
From: James Betker <jbetker@gmail.com>
Date: Sat, 5 Sep 2020 07:59:54 -0600
Subject: [PATCH] Fix srg computer bug

---
 codes/models/archs/SwitchedResidualGenerator_arch.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/codes/models/archs/SwitchedResidualGenerator_arch.py b/codes/models/archs/SwitchedResidualGenerator_arch.py
index aa0d740a..b3d62ec0 100644
--- a/codes/models/archs/SwitchedResidualGenerator_arch.py
+++ b/codes/models/archs/SwitchedResidualGenerator_arch.py
@@ -253,6 +253,8 @@ class ConfigurableSwitchComputer(nn.Module):
             else:
                 x = x1 + rand_feature
 
+        if not isinstance(x, tuple):
+            x = (x,)
         if self.pre_transform:
             x = self.pre_transform(*x)
         if not isinstance(x, tuple):