From 1ad35a3c4ae7d9743ee385cd13a90a1bcec14b92 Mon Sep 17 00:00:00 2001
From: Austin Garrett <agarret7@mit.edu>
Date: Mon, 9 Apr 2018 17:00:45 -0400
Subject: [PATCH] Fix MNIST to work again.

---
 core/util_layers.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/core/util_layers.py b/core/util_layers.py
index 684af98..5208d7c 100644
--- a/core/util_layers.py
+++ b/core/util_layers.py
@@ -10,7 +10,13 @@ def EndChannels(f, make_contiguous = False):
 
         def __init__(self):
             super(WrappedLayer, self).__init__()
-            self.forward = lambda x: f(x.permute(0,3,1,2)).permute(0,2,3,1)
+            self.f = f
+
+        def forward(self, x):
+            x = x.permute(0,3,1,2)
+            x = self.f(x)
+            x = x.permute(0,2,3,1)
+            return x
 
     return WrappedLayer()