Hello everyone, hope you are all having a great day.
I was wondering if its possible to convert the following module into torch script:
class PriorBox(torch.nn.Module):
    def __init__(self):
        super(PriorBox, self).__init__()
    
    # @torch.jit.script
    def forward(self, minSizes, steps, clip, image_size):
        anchors = []
        feature_maps = [[ceil(image_size[0]/step), ceil(image_size[1]/step)] for step in steps]

        for k, f in enumerate(feature_maps):
            min_sizes = minSizes[k]
            for i, j in product(range(f[0]), range(f[1])):
                for min_size in min_sizes:
                    s_kx = min_size / image_size[1]
                    s_ky = min_size / image_size[0]
                    dense_cx = [x * steps[k] / image_size[1] for x in [j + 0.5]]
                    dense_cy = [y * steps[k] / image_size[0] for y in [i + 0.5]]
                    for cy, cx in product(dense_cy, dense_cx):
                        anchors += [cx, cy, s_kx, s_ky]

        # back to torch land
        output = torch.Tensor(anchors).view(-1, 4)
        if clip:
            output.clamp_(max=1, min=0)
        return output

I have reimplemented this in libtorch but its slow as snail!
here it is in case its needed:
struct PriorBox : torch::nn::Module
{
    PriorBox(const std::vector<std::pair<torch::Tensor, torch::Tensor>>& min_sizes, const std::vector<int>& steps, const bool& clip, const std::pair<int, int>& image_size, const std::string& phase = "train")
    {
        this->min_sizes = min_sizes;
        this->steps = steps;
        this->clip = clip;
        this->image_size = image_size;
        this->phase = phase;
        this->name = name;
        for (auto& step : this->steps)
        {
            auto height = torch::tensor(float(this->image_size.first) / step);
            auto width = torch::tensor(float(this->image_size.second) / step);
            this->feature_maps.emplace_back(std::make_pair(torch::ceil(height), torch::ceil(width)));
        }
    }
    
    torch::Tensor forward()
    {
        std::vector<torch::Tensor> anchors;
        int i = -1;
        for (auto& fmap : this->feature_maps)
        {
            auto min_sizes = this->min_sizes[++i];
            auto result = torch::cartesian_prod({ torch::arange(c10::Scalar(0), c10::Scalar(fmap.first.item().toInt()), c10::Scalar(1)),
                                                  torch::arange(c10::Scalar(0), c10::Scalar(fmap.second.item().toInt()), c10::Scalar(1)) });
            for (int idx = 0; idx <= result.sizes()[0] - 1; idx++)
            {
                //takes around 0.006 ms
                auto i_ = result[idx][0];
                auto j_ = result[idx][1];
                //takes around 0.20 ms to 0.30 ms
                for (auto& min_size : { min_sizes.first, min_sizes.second })
                {
                    auto s_kx = min_size / float(this->image_size.second);
                    auto s_ky = min_size / float(this->image_size.first);
                    // takes around 0.037 ms
                    torch::Tensor dense_cx = (j_ + 0.5) * this->steps[i] / this->image_size.second;
                    torch::Tensor dense_cy = (i_ + 0.5) * this->steps[i] / this->image_size.first;
                    //takes around 0.02ms
                    auto result_cy_cx = torch::cartesian_prod({ dense_cy.unsqueeze(0), dense_cx.unsqueeze(0) });
                    //this takes around 0.010ms
                    for (int l = 0; l <= result_cy_cx.sizes()[0] - 1; l++)
                    {
                        auto cy = result_cy_cx[l][0].unsqueeze(0);
                        auto cx = result_cy_cx[l][1].unsqueeze(0);
                        anchors.emplace_back(torch::cat({ cx, cy, s_kx, s_ky }));
                    }
                }
            }
        }
        //takes around 5ms!
        auto output = torch::stack(anchors).view({ -1,4 });
        //takes around 0 ms!
        if (this->clip)
            output.clamp_(0, 1);
        
        return output;
    }

    std::vector<std::pair<torch::Tensor, torch::Tensor>> min_sizes;
    std::vector<int> steps;
    bool clip;
    std::pair<int, int> image_size;
    std::string phase;
    std::string name;
    std::vector<std::pair<torch::Tensor, torch::Tensor>> feature_maps;
};


So I thought maybe converting this into a Torchscript and loading it in C++ would be a better idea.
Since we are dealing with loops, etc we cant simply jit trace the module. So I guess I’m stuck with scripting.
When I tried to do scripting I kept getting different errors such asrange cannot be used as a value:
RuntimeError:
range cannot be used as a value:
  File "P:\ligen\layers\functions\prior_box.py", line 19
        for k, f in enumerate(feature_maps):
            min_sizes = minSizes[k]
            for i, j in product(range(int(f[0])), range(int(f[1]))):
                                ~~~~~~~~~~~~~~ <--- HERE
                for min_size in min_sizes:
                    s_kx = min_size / image_size[1]

So my question is, considering this, is it portable to torchscript? if so What am I missing here? how should I go about this?
If not, what are my other options?
Thanks a lot in advance