Mask changed to cuda byte tensor for cutorch api

sensibill · Aug 9, 2016 · bbec157 · bbec157
1 parent ce4f0e7
commit bbec157
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/MaskZero.lua b/MaskZero.lua
@@ -68,7 +68,7 @@ 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.CudaTensor() or torch.ByteTensor())
+   self.zeroMask = self.zeroMask or ((torch.type(rmi) == 'torch.CudaTensor') and torch.CudaByteTensor() or torch.ByteTensor())
    self._zeroMask.eq(self.zeroMask, self._zeroMask, 0)
 
    -- forward through decorated module