Merge pull request #103 from ivpopov/patch-2R3

Fixed wrong detection of tensor size 1/x/y/z = x/y/z

2 files changed, 20 insertions, 9 deletions

diff --git a/SpatialConvolution.lua b/SpatialConvolution.lua
index d691081..99d4acd 100644
--- a/SpatialConvolution.lua
+++ b/SpatialConvolution.lua
@@ -83,17 +83,19 @@ function SpatialConvolution:resetMode()
 end
 
 function SpatialConvolution:createIODescriptors(input)
-    local batch = true
-    if input:dim() == 3 then
-        input = input:view(1, input:size(1), input:size(2), input:size(3))
-        batch = false
-    end
-    assert(input:dim() == 4 and input:isContiguous());
-    self.iSize = self.iSize or torch.LongStorage(4):fill(0)
+    assert((input:dim() == 3 or input:dim() == 4) and input:isContiguous());
+
     if not self.iDesc or not self.oDesc or
-        input:size(1) ~= self.iSize[1] or input:size(2) ~= self.iSize[2]
-    or input:size(3) ~= self.iSize[3] or input:size(4) ~= self.iSize[4] then
+        input:size(1) ~= self.iSize[1] or input:size(2) ~= self.iSize[2] or
+        input:size(3) ~= self.iSize[3] or
+        (#input:size() == 4 and #self.iSize == 4 and
+         input:size(4) ~= self.iSize[4]) then
         self.iSize = input:size()
+        local batch = true
+        if input:dim() == 3 then
+            input = input:view(1, input:size(1), input:size(2), input:size(3))
+            batch = false
+        end
 
         -- resize gradInput
         if self.gradInput then self.gradInput:resizeAs(input); end
diff --git a/test/test.lua b/test/test.lua
index ea2ffce..e55e34f 100644
--- a/test/test.lua
+++ b/test/test.lua
@@ -221,6 +221,15 @@ function cudnntest.SpatialConvolution_backward_single()
   test(sconv, gconv)
 end
 
+function cudnntest.SpatialConvolution_batch_reshape()
+   local gconv = cudnn.SpatialConvolution(4, 8, 4, 4, 2, 2):cuda()
+   local outputBatch = gconv:forward(torch.rand(1, 4, 32, 32):cuda())
+   mytester:asserteq(outputBatch:dim(), 4, 'error in dimension')
+
+   local outputNoBatch = gconv:forward(torch.rand(4, 32, 32):cuda())
+   mytester:asserteq(outputNoBatch:dim(), 3, 'error in dimension')
+end
+
 function cudnntest.VolumetricConvolution_forward_single()
    local from = math.random(1,16)
    local to = math.random(1,16)