Definitive HeisenDistilBug fix

cc @julien-c @@thomwolf
2020-01-27 12:09:58 -05:00
parent f09f42d4d3
commit 875c4ae48f
2 changed files with 5 additions and 1 deletions
--- 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: