go hey

1 files changed, 6 insertions, 5 deletions

diff --git a/run.py b/run.py
index 9c0bae2..c705552 100644
--- a/run.py
+++ b/run.py
@@ -53,13 +53,14 @@ def process_tree(moduleNetwork, a, b, tensorOutput, steps):
 	tensorInputFirst = torch.FloatTensor(a)
 	tensorInputSecond = torch.FloatTensor(b)
 	process(moduleNetwork, tensorInputFirst, tensorInputSecond, tensorOutput)
-	tensorMiddle = (numpy.rollaxis(tensorOutput.clamp(0.0, 1.0).numpy(), 0, 3)[:,:,::-1] * 255.0).astype(numpy.uint8)
+	middle = tensorOutput.clamp(0.0, 1.0).numpy()
+	np_middle = (numpy.rollaxis(output, 0, 3)[:,:,::-1] * 255.0).astype(numpy.uint8)
 	if steps < 2:
-		return [tensorMiddle]
+		return [np_middle]
 	else:
-		tensorLeft = process_tree(moduleNetwork, a, tensorMiddle, tensorOutput, steps / 2)
-		tensorRight = process_tree(moduleNetwork, tensorMiddle, b, tensorOutput, steps / 2)
-		return tensorLeft + [tensorMiddle] + tensorRight
+		left = process_tree(moduleNetwork, a, middle, tensorOutput, steps / 2)
+		right = process_tree(moduleNetwork, middle, b, tensorOutput, steps / 2)
+		return left + [np_middle] + right
 
 if arguments_strVideo and arguments_strVideoOut:
 	reader = FFMPEG_VideoReader(arguments_strVideo, False)