Wrapped flag attribute as scalar tensor

[prclibo]

Hi, I was trying to integrate your code in pytorch lightning framework and found that some scalar flag attributes caused exceptions when parallelizing the nn.Module. This PR removes it for safety and should not change code behavior or performance.