Add THHalfTensor support to cutorch (#655)

* Add THHalfTensor support to cutorch.

1 files changed, 16 insertions, 2 deletions

diff --git a/init.lua b/init.lua
index 09f00e3..fdb7b08 100644
--- a/init.lua
+++ b/init.lua
@@ -16,8 +16,8 @@ torch.CudaTensor.__tostring__        = torch.FloatTensor.__tostring__
 torch.CudaDoubleStorage.__tostring__ = torch.DoubleStorage.__tostring__
 torch.CudaDoubleTensor.__tostring__  = torch.DoubleTensor.__tostring__
 if cutorch.hasHalf then
-   torch.CudaHalfStorage.__tostring__  = torch.FloatStorage.__tostring__
-   torch.CudaHalfTensor.__tostring__  = torch.FloatTensor.__tostring__
+   torch.CudaHalfStorage.__tostring__  = torch.HalfStorage.__tostring__
+   torch.CudaHalfTensor.__tostring__  = torch.HalfTensor.__tostring__
 end
 
 require('cutorch.Tensor')
@@ -57,6 +57,20 @@ function cutorch.createCudaHostTensor(...)
    return torch.FloatTensor(storage, 1, size:storage())
 end
 
+function cutorch.createCudaHostDoubleTensor(...)
+   local size = longTensorSize(...)
+   local storage = torch.DoubleStorage(cutorch.CudaHostAllocator, size:prod())
+   return torch.DoubleTensor(storage, 1, size:storage())
+end
+
+if cutorch.hasHalf then
+  function cutorch.createCudaHostHalfTensor(...)
+     local size = longTensorSize(...)
+     local storage = torch.HalfStorage(cutorch.CudaHostAllocator, size:prod())
+     return torch.HalfTensor(storage, 1, size:storage())
+   end
+ end
+
 -- Creates a CudaTensor using the CudaUVAAllocator.
 -- Accepts either a LongStorage or a sequence of numbers.
 local function _createUVATensor(...)