from torch import Tensor
import torch.nn as nn
import torch.fx as fx

class M(nn.Module):
    def forward(self, x):
        return x

class Container(nn.Module):
    def __init__(self):
        super().__init__()
        self._module = M()

    def forward(self, *args):
        return self._module(args)  # set a variable to store

model = Container()
model(Tensor())

trace = fx.symbolic_trace(model)