minor

Signed-off-by: realAsma <[email protected]>

Author: realAsma

modelopt/torch/quantization/plugins/transformers_trainer.py

@@ -298,7 +298,7 @@ def _modelopt_prepare(self, *args, **kwargs):
 
             for tq in (m for m in model.modules() if isinstance(m, TensorQuantizer)):
                 tq._non_persistent_buffers_set.clear()
-                tq._non_persistent_buffers_set = tq_og_non_prsist_buffers[tq]
+                tq._non_persistent_buffers_set.update(tq_og_non_prsist_buffers[tq])
 
             return outputs