fix depth opt arg in totensor

2 files changed, 4 insertions, 4 deletions

diff --git a/generic/qttorch.cpp b/generic/qttorch.cpp
index 0ff66e7..2968b0f 100644
--- a/generic/qttorch.cpp
+++ b/generic/qttorch.cpp
@@ -105,7 +105,7 @@ static int qttorch_(qimage_totensor)(lua_State *L)
   }
   else
   {
-    depth = luaL_optinteger(L, 2, 3);
+    depth = luaL_optinteger(L, 4, 3);
     if (depth != 1 && depth != 3 && depth != 4)
       luaL_error(L, "depth must be 1, 3, or 4.");
     if (depth == 1)
diff --git a/init.lua b/init.lua
index 4cd88ea..c3cdbf4 100644
--- a/init.lua
+++ b/init.lua
@@ -6,11 +6,11 @@ qt.QImage.fromTensor = function(tensor, scale)
                           return tensor.qttorch.QImageFromTensor(tensor, scale)
                        end
 
-qt.QImage.toTensor = function(self, tensor, scale)
+qt.QImage.toTensor = function(self, tensor, scale, depth)
                         if type(tensor) == 'userdata' then
-                           return tensor.qttorch.QImageToTensor(self, tensor, scale)
+                           return tensor.qttorch.QImageToTensor(self, tensor, scale, depth)
                         else
                            local t = torch.getmetatable(torch.getdefaulttensortype())
-                           return t.qttorch.QImageToTensor(self, tensor, scale)
+                           return t.qttorch.QImageToTensor(self, tensor, scale, depth)
                         end
                      end