NameError in Graph Mode with Dynamic Module Addition to `nn.ModuleDict` in Forward Pass

import oneflow.nn as nn
import oneflow as flow
from oneflow import nn


class InputDepthAdapter(nn.Module):

    def __init__(self, channels):
        super().__init__()
        self.channels = channels
        self.conv_layers = nn.ModuleDict()

    def adopt_depth(self, input):
        input_depth = input.shape[1]
        input_depth_name = str(input_depth)
        if not hasattr(self, input_depth_name):
            reshape_conv = nn.Conv2d(input_depth, self.channels, kernel_size=3,
                stride=1, padding=1)
            reshape_conv.to(input.device)
            self.conv_layers[input_depth_name] = reshape_conv
        reshape_conv = getattr(self.conv_layers, input_depth_name)
        return reshape_conv(input)

    def forward(self, input):
        input = self.adopt_depth(input)
        return input


class TestModel(nn.Module):

    def __init__(self):
        super().__init__()
        self.adapter = InputDepthAdapter(64)

    def forward(self, x):
        return self.adapter(x)


input_tensor = flow.randn(8, 3, 32, 32)
model = TestModel()

# Eager mode works
print("Eager mode result:", model(input_tensor))

# Graph mode fails
class TestGraph(nn.Graph):
    def __init__(self, model):
        super().__init__()
        self.model = model
        
    def build(self, x):
        return self.model(x)

g = TestGraph(model)
print("Graph mode result:", g(input_tensor))

[ERROR](GRAPH:TestGraph_0:TestGraph) building graph got error.
Traceback (most recent call last):
  File "test_script.py", line 60, in <module>
    print("Graph mode result:", g(input_tensor))
  ...
  File "test_script.py", line 20, in forward
    self.conv_layers[input_depth_name] = reshape_conv
  File "/home/miniconda3/envs/oneflow-test/lib/python3.10/site-packages/oneflow/nn/utils/container.py", line 207, in __setitem__
    self.add_module(key, module)
  File "/home/miniconda3/envs/oneflow-test/lib/python3.10/site-packages/oneflow/nn/graph/proxy.py", line 314, in add_module
    get_block_cls(module)(
NameError: name 'get_block_cls' is not defined. Did you mean: 'get_proxy_cls'?