AssertionError

[JasonYangCode]

File "C:\ProgramData\Anaconda3\lib\site-packages\torchstat\statistics.py", line 71, in stat
ms.show_report()
File "C:\ProgramData\Anaconda3\lib\site-packages\torchstat\statistics.py", line 64, in show_report
collected_nodes = self._analyze_model()
File "C:\ProgramData\Anaconda3\lib\site-packages\torchstat\statistics.py", line 57, in _analyze_model
model_hook = ModelHook(self._model, self._input_size)
File "C:\ProgramData\Anaconda3\lib\site-packages\torchstat\model_hook.py", line 24, in init
self._model(x)
File "C:\ProgramData\Anaconda3\lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "C:\ProgramData\Anaconda3\lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "C:\ProgramData\Anaconda3\lib\site-packages\torch\nn\modules\container.py", line 217, in forward
input = module(input)
File "C:\ProgramData\Anaconda3\lib\site-packages\torch\nn\modules\module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "C:\ProgramData\Anaconda3\lib\site-packages\torch\nn\modules\container.py", line 217, in forward
input = module(input)
File "C:\ProgramData\Anaconda3\lib\site-packages\torchstat\model_hook.py", line 44, in wrap_call
assert module.class in self._origin_call
AssertionError

[JasonYangCode]

I only made the following modifications to my source code:

from rbn import RepresentativeBatchNorm2d
... ...
# nn.BatchNorm2d(oup),
RepresentativeBatchNorm2d(oup),
... ...