[QUESTION] How to quantize weight of te.linear in megatron?

I want to fake quantize weight of te.linear layer in megatron. If I use pytorch, the code is simple like this:
```
def forward(self, x_quant):
w_quant = fake_quant(self.weight)
y = torch.nn.functional.linear(x_quant, w_quant)
return y
```
However, when applying this method to te.linear, it will raise an error due to `w_quant` lacking `main.grad`. How can I solve this?