fix XLM tests

pytorch_transformers/modeling_xlm.py

@@ -566,7 +566,7 @@ class XLMPredLayer(nn.Module):
             scores = self.proj(x)
             outputs = (scores,) + outputs
             if y is not None:
-                loss = F.cross_entropy(scores.view(-1, self.n_words), y, reduction='elementwise_mean')
+                loss = F.cross_entropy(scores.view(-1, self.n_words), y.view(-1), reduction='elementwise_mean')
                 outputs = (loss,) + outputs
         else:
             scores = self.proj.log_prob(x)

pytorch_transformers/tests/modeling_xlm_test.py

@@ -185,7 +185,7 @@ class XLMModelTest(CommonTestCases.CommonModelTester):
             model.eval()
 
             outputs = model(input_ids)
-            start_top_log_probs, start_top_index, end_top_log_probs, end_top_index, cls_logits, mems = outputs
+            start_top_log_probs, start_top_index, end_top_log_probs, end_top_index, cls_logits = outputs
 
             outputs = model(input_ids, start_positions=sequence_labels,
                                          end_positions=sequence_labels,