diff --git a/tests/test_modeling_tf_common.py b/tests/test_modeling_tf_common.py
index de879aa15bf..97bccd2cbe4 100644
--- a/tests/test_modeling_tf_common.py
+++ b/tests/test_modeling_tf_common.py
@@ -112,8 +112,12 @@ class TFModelTesterMixin:
             tfo = tf_model(inputs_dict, training=False)
             tf_hidden_states = tfo[0].numpy()
             pt_hidden_states = pto[0].numpy()
+
+            pt_hidden_states[np.isnan(tf_hidden_states)] = 0
             tf_hidden_states[np.isnan(tf_hidden_states)] = 0
             pt_hidden_states[np.isnan(pt_hidden_states)] = 0
+            tf_hidden_states[np.isnan(pt_hidden_states)] = 0
+
             max_diff = np.amax(np.abs(tf_hidden_states - pt_hidden_states))
             # Debug info (remove when fixed)
             if max_diff >= 2e-2:
diff --git a/tests/test_modeling_tf_distilbert.py b/tests/test_modeling_tf_distilbert.py
index 4ecbe53ede7..5546e7a5b85 100644
--- a/tests/test_modeling_tf_distilbert.py
+++ b/tests/test_modeling_tf_distilbert.py
@@ -219,5 +219,5 @@ class TFDistilBertModelTest(TFModelTesterMixin, unittest.TestCase):
     # @slow
     # def test_model_from_pretrained(self):
     #     for model_name in list(DISTILBERT_PRETRAINED_MODEL_ARCHIVE_MAP.keys())[:1]:
-    #         model = DistilBertModel.from_pretrained(model_name, cache_dir=CACHE_DIR)
+    #         model = DistilBertModesss.from_pretrained(model_name, cache_dir=CACHE_DIR)
     #         self.assertIsNotNone(model)