Add support for Half tensors to DataParallelTable.

torch · Dec 28, 2016 · 4a21b71 · 4a21b71
1 parent 4173b22
commit 4a21b71
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/DataParallelTable.lua b/DataParallelTable.lua
@@ -550,8 +550,8 @@ function DataParallelTable:_distributeTensorRecursive(dst, src, idx, n)
 
    assert(torch.isTensor(src), 'input must be a tensor or table of tensors')
    if self.typeStr == 'torch.CudaHalfTensor' then
-      assert(false,
-             'Half Tensors not supported yet by DataParallelTable')
+      assert(src:type() == self.typeStr or src:type() == 'torch.HalfTensor',
+             'input must be a CudaHalf or Half tensor')
    elseif self.typeStr == 'torch.CudaDoubleTensor' then
       assert(src:type() == self.typeStr or src:type() == 'torch.DoubleTensor',
              'input must be a CudaDouble or Double tensor')

diff --git a/test_DataParallelTable.lua b/test_DataParallelTable.lua
@@ -25,7 +25,7 @@ local t2cpu = {
 local function checkHalf()
    if cutorch.hasHalf then
        table.insert(typenames, 'torch.CudaHalfTensor')
-       t2cpu['torch.CudaHalfTensor'] = 'torch.FloatTensor'
+       t2cpu['torch.CudaHalfTensor'] = 'torch.HalfTensor'
    end
 end