From bd6f5285cd8f6a85a20f6b46bc837b5e5dffc256 Mon Sep 17 00:00:00 2001
From: Nicholas Leonard <nick@nikopia.org>
Date: Fri, 17 Oct 2014 12:52:55 -0400
Subject: [PATCH] treemax support for non-contiguous tensors

---
 treemax.lua | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/treemax.lua b/treemax.lua
index 18fc6a6..8ca9b9f 100644
--- a/treemax.lua
+++ b/treemax.lua
@@ -7,11 +7,17 @@ function torch.treemax(tensor, treeSize)
       tmb = {
          mean = tensor.new(),
          max = tensor.new(),
-         idx = torch.LongTensor()
+         idx = torch.LongTensor(),
+         copy = tensor.new()
       }
       treeMaxBuffer[torch.type(tensor)] = tmb
    end
    
+   if not tensor:isContiguous() then
+      tmb.copy:resizeAs(tensor):copy(tensor)
+      tensor = tmb.copy
+   end
+   
    local lvl = tensor
    local maxIdx, maxVal = 1, 0
    for i=1,#treeSize do