treemax support for non-contiguous tensors

nicholas-leonard · Oct 17, 2014 · bd6f528 · bd6f528
1 parent 23e645d
commit bd6f528
Showing 1 changed file with 7 additions and 1 deletion.
diff --git 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