I am trying to implement position sensitive roi pooling  (PSROIPooling)which is proposed in RFCN work.
PSROIPooling is basically ROIPooling + average pooling.
I am using the roi_pooling.py  that is written in pytorch and provided here 18.
and trying to change this part of the code 14 to be completely in pytorch (please note that the current version is in cuda, but i need to do some modification, so that is why im trying to change it to be in pytorch)
so I change that file 14 from:
import torch
from torch.autograd import Function
from .._ext import psroi_pooling 


class PSRoIPoolingFunction(Function):
    def __init__(self, pooled_height, pooled_width, spatial_scale, group_size, output_dim):
        self.pooled_width = int(pooled_width)
        self.pooled_height = int(pooled_height)
        self.spatial_scale = float(spatial_scale)
        self.group_size = int(group_size)
        self.output_dim = int(output_dim)
        self.output = None
        self.mappingchannel = None
        self.rois = None
        self.feature_size = None

    def forward(self, features, rois):
        batch_size, num_channels, data_height, data_width = features.size()
        num_rois = rois.size()[0]

        output = torch.zeros(num_rois, self.output_dim, self.pooled_height, self.pooled_width)
        mappingchannel = torch.IntTensor(num_rois, self.output_dim, self.pooled_height, self.pooled_width).zero_()
        output = output.cuda()

        mappingchannel = mappingchannel.cuda()

        psroi_pooling.psroi_pooling_forward_cuda(self.pooled_height, self.pooled_width, self.spatial_scale, self.group_size, self.output_dim, \
        features, rois, output, mappingchannel)


        
        
        self.output = output
        self.mappingchannel = mappingchannel
        self.rois = rois
        self.feature_size = features.size()

        return output

    def backward(self, grad_output):
        assert(self.feature_size is not None and grad_output.is_cuda)

        batch_size, num_channels, data_height, data_width = self.feature_size

        grad_input = torch.zeros(batch_size, num_channels, data_height, data_width).cuda()

        psroi_pooling.psroi_pooling_backward_cuda(self.pooled_height, self.pooled_width, self.spatial_scale, self.output_dim,  \
        grad_output, self.rois, grad_input, self.mappingchannel)
        return grad_input, None


to be like this:
import torch
from torch.autograd import Function
from .._ext import psroi_pooling 


from .ROI_Pooling_PyTorch import *
from .ROI_Pooling_PyTorch import roi_pooling
from torch.autograd  import Variable

class PSRoIPoolingFunction(Function):
    def __init__(self, pooled_height, pooled_width, spatial_scale, group_size, output_dim):
        self.pooled_width = int(pooled_width)
        self.pooled_height = int(pooled_height)
        self.spatial_scale = float(spatial_scale)
        self.group_size = int(group_size)
        self.output_dim = int(output_dim)
        self.output = None
        self.mappingchannel = None
        self.rois = None
        self.feature_size = None

    def forward(self, features, rois):
        batch_size, num_channels, data_height, data_width = features.size()
        num_rois = rois.size()[0]
        output = torch.zeros(num_rois, self.output_dim, self.pooled_height, self.pooled_width)
        #        mappingchannel = torch.IntTensor(num_rois, self.output_dim, self.pooled_height, self.pooled_width).zero_()

        # ROI Pooling
        out2 = roi_pooling(features, rois, size=(self.pooled_height,self.pooled_width),
                           spatial_scale = self.spatial_scale)
        
        # AVerage pooling for Position Sensitive
        
        output = Variable(output.cuda())

        chan= 0
        for i in range(0,out2.size(1),self.pooled_height*self.pooled_width):
            output[:,chan,:,:] = torch.mean(out2[:,i:i+self.pooled_height*self.pooled_width,:,:],1,keepdim=True)
            chan += 1
            
        #        mappingchannel = mappingchannel.cuda()
        
        self.output = output
        #        self.mappingchannel = mappingchannel
        self.rois = rois
        self.feature_size = features.size()
        

        return output.data

    def backward(self, grad_output):
        
# =============================================================================
#         What should i put here?????
# =============================================================================

the forward pass sounds like working, but the backward pass does not, and i have no idea what to do with it, in fact i have no idea what should i put there at all.
Can anyone pleaze help me 
ps. just for the record I am using pytorch version 0.3.1 and i cannot switch to 0.4.1 at the moment