add opencl to MaskZero.lua

torch · Aug 19, 2016 · 11e3183 · 11e3183
1 parent 88cbe99
commit 11e3183
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/.gitignore b/.gitignore
@@ -0,0 +1 @@
+build/
diff --git a/MaskZero.lua b/MaskZero.lua
@@ -68,7 +68,11 @@ function MaskZero:updateOutput(input)
    local vectorDim = rmi:dim() 
    self._zeroMask = self._zeroMask or rmi.new()
    self._zeroMask:norm(rmi, 2, vectorDim)
-   self.zeroMask = self.zeroMask or ((torch.type(rmi) == 'torch.CudaTensor') and torch.CudaByteTensor() or torch.ByteTensor())
+   self.zeroMask = self.zeroMask or (
+       (torch.type(rmi) == 'torch.CudaTensor') and torch.CudaByteTensor()
+       or (torch.type(rmi) == 'torch.ClTensor') and torch.ClTensor()
+       or torch.ByteTensor()
+    )
    self._zeroMask.eq(self.zeroMask, self._zeroMask, 0)
 
    -- forward through decorated module