Index error while running A2T, CNN, MR
**Describe the bug**
Index error while running inference of attack (A2T) over text classification model (CNN) using dataset (MR)

**To Reproduce**
Colab link: https://colab.research.google.com/drive/10pV0ArRPIG0DjmgPjIXZhldCLRAciYbP?usp=sharing


**Expected behavior**
Running smooth inference over all samples of MR dataset like another combination of attacks and classification models.

**Screenshots or Traceback**
mr__cnn__a2t -- error
textattack: Goal function <class 'textattack.goal_functions.classification.untargeted_classification.UntargetedClassification'> compatible with model WordCNNForClassification.
  0%|          | 0/1066 [00:00<?, ?it/s]/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py:1025: UserWarning: Using a non-full backward hook when the forward contains multiple autograd Nodes is deprecated and will be removed in future versions. This hook will be missing some grad_input. Please use register_full_backward_hook to get the documented behavior.
  warnings.warn("Using a non-full backward hook when the forward contains multiple autograd Nodes "
Traceback (most recent call last):
  File "<ipython-input-10-9789d5ec7ad6>", line 59, in <module>
    run_script(dataset_type, model_type, attack_type)
  File "<ipython-input-8-71af6dc8bf0c>", line 2, in run_script
    results, time_taken  = run_model(dataset_type, model_type, attack_type)
  File "<ipython-input-5-20a74e5deceb>", line 29, in run_model
    results = attacker.attack_dataset()
  File "/usr/local/lib/python3.7/dist-packages/textattack/attacker.py", line 441, in attack_dataset
    self._attack()
  File "/usr/local/lib/python3.7/dist-packages/textattack/attacker.py", line 170, in _attack
    raise e
  File "/usr/local/lib/python3.7/dist-packages/textattack/attacker.py", line 168, in _attack
    result = self.attack.attack(example, ground_truth_output)
  File "/usr/local/lib/python3.7/dist-packages/textattack/attack.py", line 422, in attack
    result = self._attack(goal_function_result)
  File "/usr/local/lib/python3.7/dist-packages/textattack/attack.py", line 371, in _attack
    final_result = self.search_method(initial_result)
  File "/usr/local/lib/python3.7/dist-packages/textattack/search_methods/search_method.py", line 36, in __call__
    result = self.perform_search(initial_result)
  File "/usr/local/lib/python3.7/dist-packages/textattack/search_methods/greedy_word_swap_wir.py", line 122, in perform_search
    index_order, search_over = self._get_index_order(attacked_text)
  File "/usr/local/lib/python3.7/dist-packages/textattack/search_methods/greedy_word_swap_wir.py", line 101, in _get_index_order
    agg_grad = np.mean(gradient[matched_tokens], axis=0)
IndexError: index 1 is out of bounds for axis 0 with size 1
  0%|          | 0/1066 [00:00<?, ?it/s]

**System Information (please complete the following information):**
Colab link is provided!

**Additional context**
N/A

diff --git a/textattack/models/wrappers/pytorch_model_wrapper.py b/textattack/models/wrappers/pytorch_model_wrapper.py
index 2eea6e33..054f8e33 100644
--- a/textattack/models/wrappers/pytorch_model_wrapper.py
+++ b/textattack/models/wrappers/pytorch_model_wrapper.py
@@ -89,7 +89,16 @@ class PyTorchModelWrapper(ModelWrapper):
         loss.backward()
 
         # grad w.r.t to word embeddings
-        grad = torch.transpose(emb_grads[0], 0, 1)[0].cpu().numpy()
+
+        # Fix for Issue #601
+
+        # Check if gradient has shape [max_sequence,1,_] ( when model input in transpose of input sequence)
+
+        if emb_grads[0].shape[1] == 1:
+            grad = torch.transpose(emb_grads[0], 0, 1)[0].cpu().numpy()
+        else:
+            # gradient has shape [1,max_sequence,_]
+            grad = emb_grads[0][0].cpu().numpy()
 
         embedding_layer.weight.requires_grad = original_state
         emb_hook.remove()
['textattack/models/wrappers/pytorch_model_wrapper.py']