Not able to run jit.script() on spatial pyramid pooling model

I tried spatial pyramid pooling (SPP) in one of my research work, which required to save the model in lite version so that it can be used by android app.

Jit Script:

from torch.utils.mobile_optimizer import optimize_for_mobile
traced_script_module = torch.jit.script(model)
traced_script_module_optimized = optimize_for_mobile(traced_script_module)
traced_script_module_optimized._save_for_lite_interpreter ("Model_PyTorchCNN_forMobileV10_Spatial.ptl")

Ref of SPP: https://github.com/revidee/pytorch-pyramid-pooling

But I am getting the following error:

RuntimeError: 

_pad(Tensor input, int[] pad, str mode="constant", float value=0.) -> (Tensor):
Expected a value of type 'float' for argument 'value' but instead found type 'int'.
:
  File "<ipython-input-9-030eeb1a6a5a>", line 47
            assert w_pad1 + w_pad2 == (w_kernel * levels[i] - previous_conv_size[1]) and h_pad1 + h_pad2 == (h_kernel * levels[i] - previous_conv_size[0])
    
            padded_input = F.pad(input=previous_conv, pad=[w_pad1, w_pad2, h_pad1, h_pad2], 
                           ~~~~~ <--- HERE
                                 mode='constant', value=0)
            if mode == "max":
'spatial_pyramid_pool' is being compiled since it was called from 'SpatialPyramidPooling.forward'
  File "<ipython-input-10-52240195b8a3>", line 18
    def forward(self, x):
        return self.spatial_pyramid_pool(x, self.levels, self.mode)
               ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE

On the following block of code:

def spatial_pyramid_pool(previous_conv, levels, mode):
        """
        Static Spatial Pyramid Pooling method, which divides the input Tensor vertically and horizontally
        (last 2 dimensions) according to each level in the given levels and pools its value according to the given mode.
        :param previous_conv input tensor of the previous convolutional layer
        :param levels defines the different divisions to be made in the width and height dimension
        :param mode defines the underlying pooling mode to be used, can either be "max" or "avg"
        :returns a tensor vector with shape [batch x 1 x n],
                                            where n: sum(filter_amount*level*level) for each level in levels
                                            which is the concentration of multi-level pooling
        """
        num_sample = previous_conv.size(0)
        previous_conv_size = [int(previous_conv.size(2)), int(previous_conv.size(3))]
        for i in range(len(levels)):
            h_kernel = int(math.ceil(previous_conv_size[0] / levels[i]))
            w_kernel = int(math.ceil(previous_conv_size[1] / levels[i]))
            w_pad1 = int(math.floor((w_kernel * levels[i] - previous_conv_size[1]) / 2))
            w_pad2 = int(math.ceil((w_kernel * levels[i] - previous_conv_size[1]) / 2))
            h_pad1 = int(math.floor((h_kernel * levels[i] - previous_conv_size[0]) / 2))
            h_pad2 = int(math.ceil((h_kernel * levels[i] - previous_conv_size[0]) / 2))
            assert w_pad1 + w_pad2 == (w_kernel * levels[i] - previous_conv_size[1]) and h_pad1 + h_pad2 == (h_kernel * levels[i] - previous_conv_size[0])

            padded_input = F.pad(input=previous_conv, pad=[w_pad1, w_pad2, h_pad1, h_pad2], 
                                 mode='constant', value=0.0)
            if mode == "max":
                pool = nn.MaxPool2d((h_kernel, w_kernel), stride=(h_kernel, w_kernel), padding=(0, 0))
            elif mode == "avg":
                pool = nn.AvgPool2d((h_kernel, w_kernel), stride=(h_kernel, w_kernel), padding=(0, 0))
            else:
                raise RuntimeError("Unknown pooling type: %s, please use \"max\" or \"avg\".")
            x = pool(padded_input)
            if i == 0:
                spp = x.view(num_sample, -1)
            else:
                spp = torch.cat((spp, x.view(num_sample, -1)), 1)

        return spp

If I change the value of “value = 0” to “value = 0.0”, I am getting the following error:

RuntimeError: 
Arguments for call are not valid.
The following variants are available:
  
  aten::eq.Tensor(Tensor self, Tensor other) -> (Tensor):
  Expected a value of type 'Tensor' for argument 'other' but instead found type 'str'.

@ptrblck rblc
Any suggestion?