ModuleList of Sequentials doesn't work


(LSM) #1

Hi, I’m trying to export a ModuleList of Sequentials. It’s failing with the following error:

RuntimeError: 
could not export python function call <python_value>. Remove calls to python functions before export.:
@torch.jit.script_method
def forward(self, x):
    # type: (List[Tensor]) -> List[Tensor]
    outputs = []
    for cur_xc in self.xc:
        out = cur_xc(torch.tensor(0))
              ~~~~~~ <--- HERE
        outputs.append(out)
    return outputs

Example:

from fractions import gcd

import torch
from torch import nn


class CrossScale(torch.jit.ScriptModule):
    __constants__ = ['xc']

    def __init__(self, n, ng=32):
        super(CrossScale, self).__init__()

        xc = []
        for i in range(len(n)):
            m = nn.Sequential(
                nn.Conv2d(n[i], n[i], 1, bias=False),
                nn.GroupNorm(gcd(ng, n[i]), n[i]))
            xc.append(m)
        self.xc = nn.ModuleList(xc)

    @torch.jit.script_method
    def forward(self, x):
        # type: (List[Tensor]) -> List[Tensor]
        outputs = []
        for cur_xc in self.xc:
            out = cur_xc(torch.tensor(0))
            outputs.append(out)
        return outputs


if __name__ == "__main__":
    n = [32, 64, 96]
    cs = CrossScale(n)

    cs.save("cs.pt")

(Michael Suo) #2

Hi, try including the ModuleList in the __constants__ list. We have an issue up to improve the error message in this case: https://github.com/pytorch/pytorch/issues/16400


(LSM) #3

Hi, you can see in the snippet I included that the module list is already in the constants list.


(Michael Suo) #4

lol you’re right, I wasn’t reading carefully enough! Thanks for filing the GH issue, will follow up there