From 9ae850770a277161d30d515a6f26024a61e3d7c7 Mon Sep 17 00:00:00 2001
From: Plutonium-239 <plutonium.239.811@gmail.com>
Date: Thu, 22 Aug 2024 13:06:15 +0530
Subject: [PATCH] minor fixes

(cherry picked from commit c5b77165799edb097e572ef218e76e1a9f68b043)
---
 experiments/util/models.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/experiments/util/models.py b/experiments/util/models.py
index 223d2c7..7adc69a 100644
--- a/experiments/util/models.py
+++ b/experiments/util/models.py
@@ -26,9 +26,8 @@
     AutoModelForSeq2SeqLM,
     BartForConditionalGeneration,
 )
-from transformers import (
-    logging as tf_logging,
-)
+from transformers import logging as tf_logging
+from transformers import utils as tf_utils
 
 from memsave_torch.nn import (
     MemSaveConv2d,
@@ -349,6 +348,8 @@ def __init__(
 
 tf_logging.disable_progress_bar()
 tf_logging.set_verbosity_error()
+tf_utils.logging.captureWarnings(True)
+
 
 hf_transformers_models_map = {
     "gpt2": _HF_model("gpt2", {}, lm_head_name="lm_head"),