Add THNN conversion of {Spatial(AdaptiveMax,Average,Max)Pooling}

1 files changed, 19 insertions, 2 deletions

diff --git a/SpatialAveragePooling.lua b/SpatialAveragePooling.lua
index d8ef41f..231177d 100644
--- a/SpatialAveragePooling.lua
+++ b/SpatialAveragePooling.lua
@@ -45,7 +45,15 @@ end
 
 function SpatialAveragePooling:updateOutput(input)
    backwardCompatible(self)
-   input.nn.SpatialAveragePooling_updateOutput(self, input)
+   input.THNN.SpatialAveragePooling_updateOutput(
+      input:cdata(),
+      self.output:cdata(),
+      self.kW, self.kH,
+      self.dW, self.dH,
+      self.padW, self.padH,
+      self.ceil_mode,
+      self.count_include_pad
+   )
    -- for backward compatibility with saved models
    -- which are not supposed to have "divide" field
    if not self.divide then
@@ -56,7 +64,16 @@ end
 
 function SpatialAveragePooling:updateGradInput(input, gradOutput)
    if self.gradInput then
-      input.nn.SpatialAveragePooling_updateGradInput(self, input, gradOutput)
+      input.THNN.SpatialAveragePooling_updateGradInput(
+         input:cdata(),
+         gradOutput:cdata(),
+         self.gradInput:cdata(),
+         self.kW, self.kH,
+         self.dW, self.dH,
+         self.padW, self.padH,
+         self.ceil_mode,
+         self.count_include_pad
+      )
       -- for backward compatibility
       if not self.divide then
          self.gradInput:mul(self.kW*self.kH)