Export to JIT script

[OnceJune]

Hi, I tried to export the model to JIT script but got this error:
reversed(Tensor 0) -> (Tensor 0):
Expected a value of type 'Tensor' for argument '0' but instead found type 'torch.torch.nn.modules.container.ModuleList
Seems JIT doesn't not support reversed, how to solve this?

[p0p4k]

Does this help?

[OnceJune]

seems not the same issue. you can try below scripts:

import torch
import torch.nn as nn
import torch.jit as jit

class ReverseModule(nn.Module):
    def __init__(self):
        super(ReverseModule, self).__init__()
        self.layers = nn.ModuleList([nn.Linear(10, 10) for _ in range(5)])

    def forward(self, x):
        reversed_layers = list(reversed(self.layers))
        for layers in reversed_layers:
            x = layer(x)
        return x

model = ReverseModule()
scripted_module = torch.jit.script(model)

Where the reversed_layers = list(reversed(self.layers)) is the same as

vits2_pytorch/models.py

Line 122 in 379527d

flows = list(reversed(self.flows))

Torch.jit does not have support for reversed()

[p0p4k]

Will something like this work for flows?

import torch
import torch.nn as nn
import torch.jit as jit

class ReverseModule(nn.Module):
    def __init__(self):
        super(ReverseModule, self).__init__()
        self.layers = nn.ModuleList([nn.Linear(10, 10) for _ in range(5)])
        # convert to a mutable dict
        self.layers = nn.ModuleDict({str(i): layer for i, layer in enumerate(self.layers)})
        self.reversed_layers = nn.ModuleDict()
        for i in range(len(self.layers) - 1, -1, -1):
            self.reversed_layers[str(i)] = self.layers[str(i)]

    def forward(self, x):
        for k, v in self.reversed_layers.items():
            x = v(x)
        return x

model = ReverseModule()
scripted_module = torch.jit.script(model)

[EmreOzkose]

Hi @OnceJune , what are your versions of pytorch and repo ? I am trying to export jit and faces lots of issues like typing, missing module issues, etc..