and, this is the network that I am trying to trace:
import torch
import torch.nn as nn
import torch.nn.functional as F


class RPN(nn.Module):
    def __init__(self):
        super(RPN, self).__init__()

    def forward(self, z_f, x_f):
        raise NotImplementedError


def conv2d_group(x, kernel):
    batch = kernel.size()[0]
    pk = kernel.view(-1, x.size()[1], kernel.size()[2], kernel.size()[3])
    px = x.view(1, -1, x.size()[2], x.size()[3])
    po = F.conv2d(px, pk, groups=batch)
    po = po.view(batch, -1, po.size()[2], po.size()[3])
    return po


class UPChannelRPN(RPN):
    def __init__(self, anchor_num=5, feature_in=256, feature_out=256):
        super(UPChannelRPN, self).__init__()

        self.anchor_num = anchor_num
        self.feature_in = feature_in
        self.feature_out = feature_out

        self.cls_output = 2 * self.anchor_num
        self.loc_output = 4 * self.anchor_num

        self.template_cls_conv = nn.Conv2d(self.feature_in, self.feature_out * self.cls_output, kernel_size=3)
        self.template_loc_conv = nn.Conv2d(self.feature_in, self.feature_out * self.loc_output, kernel_size=3)

        self.search_cls_conv = nn.Conv2d(self.feature_in, self.feature_out, kernel_size=3)
        self.search_loc_conv = nn.Conv2d(self.feature_in, self.feature_out, kernel_size=3)

        self.loc_adjust = nn.Conv2d(self.loc_output, self.loc_output, kernel_size=1)

    def forward(self, z_f, x_f):
        cls_kernel = self.template_cls_conv(z_f)
        loc_kernel = self.template_loc_conv(z_f)

        cls_feature = self.search_cls_conv(x_f)
        loc_feature = self.search_loc_conv(x_f)

        pred_cls = conv2d_group(cls_feature, cls_kernel)
        pred_loc = self.loc_adjust(conv2d_group(loc_feature, loc_kernel))
        return pred_cls, pred_loc

I am trying to trace UPChannelRPN(). This network does not include any conditionals (to the best of my knowledge) and therefore I think it should be possible to trace this network using torch.jit.trace