@@ -8,6 +8,7 @@
 
 import lru
 import torch
+import transformers
 
 import textattack
 from textattack.attack_results import (
@@ -212,13 +213,23 @@ def cuda_(self):
         def to_cuda(obj):
             visited.add(id(obj))
             if isinstance(obj, torch.nn.Module):
-                if getattr(obj, "hf_device_map", None):
+                if isinstance(obj, transformers.PreTrainedModel) and getattr(
+                    obj, "hf_device_map", None
+                ):
                     # Model was already loaded with `device_map=...` (e.g.
                     # split across multiple GPUs via `accelerate`);
                     # `hf_device_map` is the marker attribute HuggingFace
                     # sets in that case. Forcing it onto a single device
                     # here breaks that placement. See
                     # https://github.com/QData/TextAttack/issues/798
+                    #
+                    # Scoped to `transformers.PreTrainedModel` specifically
+                    # (rather than any `torch.nn.Module`) since this visitor
+                    # also traverses non-HuggingFace models (TensorFlow,
+                    # sklearn, custom PyTorch modules, ...) reachable from
+                    # a Constraint/GoalFunction/Transformation, and only
+                    # `transformers`-loaded models actually get this
+                    # attribute set by `from_pretrained(device_map=...)`.
                     pass
                 else:
                     obj.to(textattack.shared.utils.device)
 
@@ -17,6 +17,7 @@
 
 from textattack.goal_function_results import GoalFunctionResultStatus
 from textattack.search_methods import SearchMethod
+from textattack.shared import AttackedText
 from textattack.shared.validators import (
     transformation_consists_of_word_swaps_and_deletions,
 )
@@ -112,9 +113,38 @@ def _get_index_order(self, initial_text, max_len=-1):
         elif self.wir_method == "gradient":
             victim_model = self.get_victim_model()
             index_scores = np.zeros(len_text)
-            grad_output = victim_model.get_grad(initial_text.tokenizer_input)
+            gradient_text = initial_text
+            if (
+                self.truncate_words_to is not None
+                and indices_to_order
+                and isinstance(initial_text.tokenizer_input, str)
+            ):
+                # `truncate_words_to` above only shortens the cheap post-hoc
+                # index-scoring loop; the actual expensive step is
+                # `get_grad`'s tokenize + forward + backward pass, which
+                # otherwise still runs over the full untruncated text. Bound
+                # it too by feeding it only the same word span
+                # `indices_to_order` was already limited to. Build a
+                # separate `AttackedText` for this (rather than just slicing
+                # the string passed to `get_grad`) and use it for the
+                # `align_with_model_tokens` call below too, so the returned
+                # `gradient` array and `word2token_mapping`'s token indices
+                # stay consistent with each other.
+                #
+                # Only done for single-sequence inputs: for paired inputs
+                # (e.g. premise/hypothesis), `tokenizer_input` is a tuple and
+                # truncating it here would need to preserve that pair
+                # structure (and the per-segment word budget) to avoid
+                # breaking the tokenizer's dual-sequence encoding, which is
+                # out of scope here. Those still fall back on the
+                # tokenizer's own `model_max_length` truncation inside
+                # `get_grad` as a (looser) bound.
+                last_index = max(indices_to_order)
+                truncated_str = initial_text.text_window_around_index(0, last_index + 1)
+                gradient_text = AttackedText(truncated_str)
+            grad_output = victim_model.get_grad(gradient_text.tokenizer_input)
             gradient = grad_output["gradient"]
-            word2token_mapping = initial_text.align_with_model_tokens(victim_model)
+            word2token_mapping = gradient_text.align_with_model_tokens(victim_model)
             for i, index in enumerate(indices_to_order):
                 matched_tokens = word2token_mapping[index]
                 if not matched_tokens: