Adapt nn code to changes in THNN and THCUNN

11 files changed, 16 insertions, 9 deletions

diff --git a/torch/_thnn/__init__.py b/torch/_thnn/__init__.py
index a3bc6f4c4d..97474692eb 100644
--- a/torch/_thnn/__init__.py
+++ b/torch/_thnn/__init__.py
@@ -46,8 +46,6 @@ type2backend = Backends()
 
 _thnn_headers = parse_header(THNN_H_PATH)
 _thcunn_headers = parse_header(THCUNN_H_PATH)
-for function in _thcunn_headers:
-    function.name = function.name[4:]
 
 for t in ['Float', 'Double']:
     backend = Backend(t, 'torch._thnn._THNN', _thnn_headers)
diff --git a/torch/_thnn/utils.py b/torch/_thnn/utils.py
index c1daebe7ed..c62fc2a29e 100644
--- a/torch/_thnn/utils.py
+++ b/torch/_thnn/utils.py
@@ -90,7 +90,7 @@ def parse_header(path):
                 fn_name = fn_name[:-1]
             generic_functions.append(Function(fn_name))
         elif l:
-            t, name = l.split(' ')
+            t, name = l.split()
             if '*' in name:
                 t = t + '*'
                 name = name[1:]
diff --git a/torch/csrc/cuda/Tensor.h b/torch/csrc/cuda/Tensor.h
index 7320c09bdb..4a0b20400c 100644
--- a/torch/csrc/cuda/Tensor.h
+++ b/torch/csrc/cuda/Tensor.h
@@ -8,7 +8,7 @@
 
 #define THCPDoubleTensor_Check(obj)  PyObject_IsInstance(obj, THCPDoubleTensorClass)
 #define THCPFloatTensor_Check(obj)   PyObject_IsInstance(obj, THCPFloatTensorClass)
-#define THCPHalfTensor_Check(obj)   PyObject_IsInstance(obj, THCPHalfTensorClass)
+#define THCPHalfTensor_Check(obj)    PyObject_IsInstance(obj, THCPHalfTensorClass)
 #define THCPLongTensor_Check(obj)    PyObject_IsInstance(obj, THCPLongTensorClass)
 #define THCPIntTensor_Check(obj)     PyObject_IsInstance(obj, THCPIntTensorClass)
 #define THCPShortTensor_Check(obj)   PyObject_IsInstance(obj, THCPShortTensorClass)
@@ -17,6 +17,7 @@
 
 #define THCPDoubleTensor_CData(obj)  (obj)->cdata
 #define THCPFloatTensor_CData(obj)   (obj)->cdata
+#define THCPHalfTensor_CData(obj)    (obj)->cdata
 #define THCPLongTensor_CData(obj)    (obj)->cdata
 #define THCPIntTensor_CData(obj)     (obj)->cdata
 #define THCPShortTensor_CData(obj)   (obj)->cdata
diff --git a/torch/legacy/nn/MultiLabelMarginCriterion.py b/torch/legacy/nn/MultiLabelMarginCriterion.py
index 2f143f073b..037ad013e5 100644
--- a/torch/legacy/nn/MultiLabelMarginCriterion.py
+++ b/torch/legacy/nn/MultiLabelMarginCriterion.py
@@ -11,6 +11,7 @@ class MultiLabelMarginCriterion(Criterion):
 
     def updateOutput(self, input, target):
         self.output_tensor = self.output_tensor or input.new(1)
+        target = target.long()
         self._backend.MultiLabelMarginCriterion_updateOutput(
             self._backend.library_state,
             input,
@@ -23,6 +24,7 @@ class MultiLabelMarginCriterion(Criterion):
         return self.output
 
     def updateGradInput(self, input, target):
+        target = target.long()
         self._backend.MultiLabelMarginCriterion_updateGradInput(
             self._backend.library_state,
             input,
diff --git a/torch/legacy/nn/MultiMarginCriterion.py b/torch/legacy/nn/MultiMarginCriterion.py
index fca03a75e9..213f36931c 100644
--- a/torch/legacy/nn/MultiMarginCriterion.py
+++ b/torch/legacy/nn/MultiMarginCriterion.py
@@ -17,6 +17,7 @@ class MultiMarginCriterion(Criterion):
 
     def updateOutput(self, input, target):
         self.output_tensor = self.output_tensor or input.new(1)
+        target = target.long()
         self._backend.MultiMarginCriterion_updateOutput(
             self._backend.library_state,
             input,
@@ -32,6 +33,7 @@ class MultiMarginCriterion(Criterion):
 
 
     def updateGradInput(self, input, target):
+        target = target.long()
         self._backend.MultiMarginCriterion_updateGradInput(
             self._backend.library_state,
             input,
diff --git a/torch/legacy/nn/SpatialAdaptiveMaxPooling.py b/torch/legacy/nn/SpatialAdaptiveMaxPooling.py
index 96e127ff32..9d1adae8b4 100644
--- a/torch/legacy/nn/SpatialAdaptiveMaxPooling.py
+++ b/torch/legacy/nn/SpatialAdaptiveMaxPooling.py
@@ -12,6 +12,7 @@ class SpatialAdaptiveMaxPooling(Module):
 
     def updateOutput(self, input):
         self.indices = self.indices or input.new()
+        self.indices = self.indices.long()
         self._backend.SpatialAdaptiveMaxPooling_updateOutput(
             self._backend.library_state,
             input,
diff --git a/torch/legacy/nn/SpatialFractionalMaxPooling.py b/torch/legacy/nn/SpatialFractionalMaxPooling.py
index 705cb6fd5e..9fe7cc5907 100644
--- a/torch/legacy/nn/SpatialFractionalMaxPooling.py
+++ b/torch/legacy/nn/SpatialFractionalMaxPooling.py
@@ -93,6 +93,7 @@ class SpatialFractionalMaxPooling(Module):
 
     def updateOutput(self, input):
         self.indices = self.indices or input.new()
+        self.indices = self.indices.long()
         self._initSampleBuffer(input)
         outW, outH = self._getOutputSizes(input)
 
diff --git a/torch/legacy/nn/SpatialMaxPooling.py b/torch/legacy/nn/SpatialMaxPooling.py
index b45151b948..4a8a00e946 100644
--- a/torch/legacy/nn/SpatialMaxPooling.py
+++ b/torch/legacy/nn/SpatialMaxPooling.py
@@ -19,7 +19,7 @@ class SpatialMaxPooling(Module):
         self.padH = padH
 
         self.ceil_mode = False
-        self.indices = torch.Tensor()
+        self.indices = torch.LongTensor()
 
     def ceil(self):
         self.ceil_mode = True
@@ -31,6 +31,7 @@ class SpatialMaxPooling(Module):
 
     def updateOutput(self, input):
         self.indices = self.indices or input.new()
+        self.indices = self.indices.long()
 
         dims = input.dim()
         self.iheight = input.size(dims-2)
diff --git a/torch/legacy/nn/VolumetricMaxPooling.py b/torch/legacy/nn/VolumetricMaxPooling.py
index a9dc6b1e65..b8fcd40d63 100644
--- a/torch/legacy/nn/VolumetricMaxPooling.py
+++ b/torch/legacy/nn/VolumetricMaxPooling.py
@@ -19,7 +19,7 @@ class VolumetricMaxPooling(Module):
         self.padH = padH
 
         self.ceil_mode = False
-        self.indices = torch.Tensor()
+        self.indices = torch.LongTensor()
 
     def ceil(self):
          self.ceil_mode = True
@@ -36,6 +36,7 @@ class VolumetricMaxPooling(Module):
         self.iwidth = input.size(dims-1)
 
         self.indices = self.indices or input.new()
+        self.indices = self.indices.long()
         self._backend.VolumetricMaxPooling_updateOutput(
             self._backend.library_state,
             input,
diff --git a/torch/lib/build_all.sh b/torch/lib/build_all.sh
index 2235be44b3..3df3f0f06d 100755
--- a/torch/lib/build_all.sh
+++ b/torch/lib/build_all.sh
@@ -67,6 +67,6 @@ build libshm
 
 cp $INSTALL_DIR/lib/* .
 cp THNN/generic/THNN.h .
-cp THCUNN/THCUNN.h .
+cp THCUNN/generic/THCUNN.h .
 cp -r tmp_install/include .
 cp $INSTALL_DIR/bin/* .
diff --git a/torch/nn/functions/thnn/pooling.py b/torch/nn/functions/thnn/pooling.py
index 17ee2e7b9d..592a25cc78 100644
--- a/torch/nn/functions/thnn/pooling.py
+++ b/torch/nn/functions/thnn/pooling.py
@@ -18,7 +18,7 @@ class _MaxPoolingBase(Function):
 
     def forward(self, input):
         self._backend = type2backend[type(input)]
-        indices, output = input.new(), input.new()
+        indices, output = input.new().long(), input.new()
         self._fw_call(self._backend.library_state, input, output, indices,
                 *self.additional_args)
         if self.save_indices:
@@ -155,7 +155,7 @@ class FractionalMaxPool2d(Function):
             self.ow = int(input.size(3) * self.rw)
         assert isinstance(self.oh, int) and isinstance(self.ow, int)
 
-        indices = input.new()
+        indices = input.new().long()
         output = input.new()
         self._backend = type2backend[type(input)]
         self._backend.SpatialFractionalMaxPooling_updateOutput(