From ac99217e92c43066af7ec96554054d75532565d7 Mon Sep 17 00:00:00 2001
From: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
Date: Wed, 10 Jun 2020 09:26:06 -0400
Subject: [PATCH] Fix the CI (#4903)

* Fix CI
---
 tests/test_modeling_common.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/test_modeling_common.py b/tests/test_modeling_common.py
index b2c11bb44a..ef00cd8694 100644
--- a/tests/test_modeling_common.py
+++ b/tests/test_modeling_common.py
@@ -67,6 +67,8 @@ class ModelTesterMixin:
         if model_class in MODEL_FOR_MULTIPLE_CHOICE_MAPPING.values():
             return {
                 k: v.unsqueeze(1).expand(-1, self.model_tester.num_choices, -1).contiguous()
+                if isinstance(v, torch.Tensor) and v.ndim != 0
+                else v
                 for k, v in inputs_dict.items()
             }
         return inputs_dict
@@ -157,7 +159,7 @@ class ModelTesterMixin:
             model.to(torch_device)
             model.eval()
             with torch.no_grad():
-                outputs = model(**inputs_dict)
+                outputs = model(**self._prepare_for_class(inputs_dict, model_class))
             attentions = outputs[-1]
             self.assertEqual(model.config.output_hidden_states, False)
             self.assertEqual(len(attentions), self.model_tester.num_hidden_layers)