From 90dbcae79008d432e22aeccbfd2a2ae6c1a8b32f Mon Sep 17 00:00:00 2001
From: Konstantin <konstantin@neuralmagic.com>
Date: Tue, 8 Mar 2022 08:03:16 -0500
Subject: [PATCH 1/2] Update: datasets update to 18.4

---
 src/sparseml/transformers/__init__.py                          | 2 +-
 src/sparseml/transformers/sparsification/question_answering.py | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/sparseml/transformers/__init__.py b/src/sparseml/transformers/__init__.py
index 6a887388799..da1c516ace1 100644
--- a/src/sparseml/transformers/__init__.py
+++ b/src/sparseml/transformers/__init__.py
@@ -60,7 +60,7 @@ def _install_transformers_and_deps():
                 "pip",
                 "install",
                 transformers_requirement,
-                "datasets<1.18.0",
+                "datasets<=1.18.4",
                 "sklearn",
                 "seqeval",
             ]
diff --git a/src/sparseml/transformers/sparsification/question_answering.py b/src/sparseml/transformers/sparsification/question_answering.py
index c13a0037978..ba445e5060c 100644
--- a/src/sparseml/transformers/sparsification/question_answering.py
+++ b/src/sparseml/transformers/sparsification/question_answering.py
@@ -299,7 +299,9 @@ def postprocess_qa_predictions(
                         start_index >= len(offset_mapping)
                         or end_index >= len(offset_mapping)
                         or offset_mapping[start_index] is None
+                        or len(offset_mapping[start_index]) < 2
                         or offset_mapping[end_index] is None
+                        or len(offset_mapping[end_index]) < 2
                     ):
                         continue
                     # Don't consider answers with a length that is either < 0 or >

From b79639817aa287cfbab2662748df9819f95d3dfc Mon Sep 17 00:00:00 2001
From: Konstantin <konstantin@neuralmagic.com>
Date: Mon, 21 Mar 2022 09:31:44 -0400
Subject: [PATCH 2/2] Fix: grad scaler correct passing when no scaler

---
 src/sparseml/transformers/sparsification/trainer.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/sparseml/transformers/sparsification/trainer.py b/src/sparseml/transformers/sparsification/trainer.py
index 600c9d2be11..f88db3bfc25 100644
--- a/src/sparseml/transformers/sparsification/trainer.py
+++ b/src/sparseml/transformers/sparsification/trainer.py
@@ -251,13 +251,14 @@ def create_optimizer(self):
                 self.manager,
                 steps_per_epoch=self.manager_steps_per_epoch,
                 loggers=self.logger_manager,
-                grad_sampler=self.grad_sampler,
+                initialize_kwargs={"grad_sampler": self.grad_sampler},
             )
             if not self.manager.initialized:
                 self.manager.initialize(
                     self.model,
                     loggers=self.logger_manager,
                     distillation_teacher=self.teacher,
+                    grad_sampler=self.grad_sampler,
                 )
         self.manager_initialized = True
         _LOGGER.info(