Some doubts for total loss?

[wangyanhao0517]

Sorry to bother you. I found that your code is different from your paper in the loss calculation process. Why is the final loss NLL + beta * KL? Your paper says KL is a part of noise

[seanie12]

Gradient of KL with respect to \theta (the parameter of BERT) is zero since we do not backpropagate it to BERT (torch.no_grad() in [line 42](https://github.com/seanie12/SWEP/blob/main/models.py#L43).

So the implementation is consistent with our equation.

[wangyanhao0517]

Gradient of KL with respect to \theta (the parameter of BERT) is zero since we do not backpropagate it to BERT (torch.no_grad() in line 42.
So the implementation is consistent with our equation.
Thanks for your reply !