Merge pull request #327 from hughperkins/add-opencl

Add opencl
Element-Research · Sep 4, 2016 · 4f8401d · 4f8401d
2 parents 1cb5b3c + 8d0c555
commit 4f8401d
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 2 deletions.
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

diff --git a/MaskZeroCriterion.lua b/MaskZeroCriterion.lua
@@ -53,7 +53,8 @@ function MaskZeroCriterion:updateOutput(input, target)
    self._zeroMask = self._zeroMask or rmi.new()
    self._zeroMask:norm(rmi, 2, vectorDim)
    local zeroMask = self._zeroMask
-   if torch.isTypeOf(zeroMask, 'torch.CudaTensor') then
+   if torch.isTypeOf(zeroMask, 'torch.CudaTensor') or
+         torch.isTypeOf(zeroMask, 'torch.ClTensor') then
       self.__zeroMask = self.__zeroMask or torch.FloatTensor()
       self.__zeroMask:resize(self._zeroMask:size()):copy(self._zeroMask)
       zeroMask = self._zeroMask