TDA4VM: QuantTrainModule Error

Hi,

I have a model that has a module of nn.Sequential(nn.Linear, nn.BatchNorm1D, nn.ReLU).

When I tried to train QuantTrainModule for the model, I got an error like below.

 
 File "/home/jaehyung/src/pytorch-jacinto-ai-devkit/modules/pytorch_jacinto_ai/xnn/quantize/quant_base_module.py", line 66, in _forward_output_activation
    outputs = op.activation_q(outputs)
  File "/home/jaehyung/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/jaehyung/src/pytorch-jacinto-ai-devkit/modules/pytorch_jacinto_ai/xnn/quantize/quant_train_module.py", line 321, in forward
    conv, weight, bias = self.merge_quantize_weights(qparams, conv, bn)
  File "/home/jaehyung/src/pytorch-jacinto-ai-devkit/modules/pytorch_jacinto_ai/xnn/quantize/quant_train_module.py", line 396, in merge_quantize_weights
    merged_bias = conv.bias if (conv.bias is not None) else torch.zeros(conv.out_channels).to(conv.weight.device)
  File "/home/jaehyung/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 779, in __getattr__
    type(self).__name__, name))
torch.nn.modules.module.ModuleAttributeError: 'QuantTrainLinear' object has no attribute 'out_channels'

How can I solve the issue?

Thanks,