Export to JIT script about vits2_pytorch HOT 4 OPEN

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)

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)